Support group by more than one keys

Signed-off-by: wang.zhen <wangzhenaaa7@gmail.com>
databendlabs · wangzhen11aaa · Feb 15, 2022 · Feb 15, 2022 · Feb 15, 2022 · Feb 15, 2022
commit 581e339a03d7674beb09143c77f8063e7c464007
diff --git a/common/datablocks/src/kernels/data_block_group_by.rs b/common/datablocks/src/kernels/data_block_group_by.rs
@@ -32,14 +32,14 @@ impl DataBlock {
         let mut group_key_len = 0;
         for col in column_names {
             let column = block.try_column_by_name(col)?;
-            let _typ = if column.is_nullable() {
+            let typ = if column.is_nullable() {
                 remove_nullable(&column.data_type())
             } else {
                 column.data_type()
             };
-            if _typ.data_type_id().is_integer() {
+            if typ.data_type_id().is_integer() {
                 // If column is nullable, we will add one byte to identify `null` value for every row.
-                group_key_len += _typ.data_type_id().numeric_byte_size()?;
+                group_key_len += typ.data_type_id().numeric_byte_size()?;
                 if column.is_nullable() {
                     group_key_len += 1;
                 }

diff --git a/common/datablocks/src/kernels/data_block_group_by_hash.rs b/common/datablocks/src/kernels/data_block_group_by_hash.rs
@@ -251,32 +251,31 @@ where T: PrimitiveType
 
         let mut res = Vec::with_capacity(group_fields.len());
         let mut offsize = 0;
-        let mut null_part_offset = 0;
 
-        init_nullable_offset_via_fields(&mut null_part_offset, group_fields)?;
+        let mut null_part_offset = init_nullable_offset_via_fields(group_fields)?;
         for f in group_fields.iter() {
             let data_type = f.data_type();
             let mut deserializer = data_type.create_deserializer(rows);
             let reader = vec8.as_slice();
             if !data_type.is_nullable() {
                 deserializer.de_batch(&reader[offsize..], step, rows)?;
                 res.push(deserializer.finish_to_column());
+                offsize += data_type.data_type_id().numeric_byte_size()?;
             } else {
+                let nullable_type_size = remove_nullable(data_type)
+                    .data_type_id()
+                    .numeric_byte_size()?;
                 deserializer.de_batch_with_nullable(
                     &reader[offsize..],
                     step,
                     rows,
-                    null_part_offset,
+                    &mut null_part_offset,
+                    nullable_type_size,
                 )?;
                 res.push(deserializer.finish_to_column());
                 null_part_offset += 1;
-            }
-            if data_type.is_nullable() {
-                offsize += remove_nullable(data_type)
-                    .data_type_id()
-                    .numeric_byte_size()?;
-            } else {
-                offsize += data_type.data_type_id().numeric_byte_size()?;
+
+                offsize += nullable_type_size;
             }
         }
         Ok(res)
@@ -304,8 +303,7 @@ where
 
         let group_columns_has_nullable_one = check_group_columns_has_nullable(group_columns);
         if group_columns_has_nullable_one {
-            let mut null_part_offset = 0;
-            init_nullable_offset(&mut null_part_offset, group_columns)?;
+            let mut null_part_offset = init_nullable_offset(group_columns)?;
             while size > 0 {
                 build_keys_with_nullable_column(
                     size,
@@ -384,30 +382,29 @@ where
 
 /// Init the nullable part's offset in bytes. It follows the values part.
 #[inline]
-fn init_nullable_offset(null_part_offset: &mut usize, group_keys: &[&ColumnRef]) -> Result<()> {
+fn init_nullable_offset(group_keys: &[&ColumnRef]) -> Result<usize> {
+    let mut _null_part_offset = 0;
     for group_key_column in group_keys {
-        *null_part_offset += remove_nullable(&group_key_column.data_type())
+        _null_part_offset += remove_nullable(&group_key_column.data_type())
             .data_type_id()
             .numeric_byte_size()?;
     }
-    Ok(())
+    Ok(_null_part_offset)
 }
 
 /// Init the nullable part's offset in bytes. It follows the values part.
 #[inline]
-fn init_nullable_offset_via_fields(
-    null_part_offset: &mut usize,
-    group_fields: &[DataField],
-) -> Result<()> {
+fn init_nullable_offset_via_fields(group_fields: &[DataField]) -> Result<usize> {
+    let mut null_part_offset = 0;
     for f in group_fields {
         let f_typ = f.data_type();
         if f_typ.is_nullable() {
-            *null_part_offset += remove_nullable(f_typ).data_type_id().numeric_byte_size()?;
+            null_part_offset += remove_nullable(f_typ).data_type_id().numeric_byte_size()?;
         } else {
-            *null_part_offset += f_typ.data_type_id().numeric_byte_size()?;
+            null_part_offset += f_typ.data_type_id().numeric_byte_size()?;
         }
     }
-    Ok(())
+    Ok(null_part_offset)
 }
 
 #[inline]

diff --git a/common/datavalues/src/types/deserializations/mod.rs b/common/datavalues/src/types/deserializations/mod.rs
@@ -42,7 +42,8 @@ pub trait TypeDeserializer: Send + Sync {
         _reader: &[u8],
         _step: usize,
         _rows: usize,
-        _null_offset: usize,
+        _null_offset: &mut usize,
+        _type_size: usize,
     ) -> Result<()> {
         Err(ErrorCode::BadDataValueType(
             "de_batch_with_nullable operation only for nullable",

diff --git a/common/datavalues/src/types/deserializations/nullable.rs b/common/datavalues/src/types/deserializations/nullable.rs
@@ -52,20 +52,15 @@ impl TypeDeserializer for NullableDeserializer {
         reader: &[u8],
         step: usize,
         rows: usize,
-        null_offset: usize,
+        null_offset: &mut usize,
+        type_size: usize,
     ) -> Result<()> {
         for row in 0..rows {
             let mut reader = &reader[step * row..];
             self.inner.de(&mut reader)?;
-            // null_offset -= self.inner().to_physical_type;
-            // todo, We should get the length of missed bytes which caused
-            // by self.inner.de method.
-            if reader[null_offset - 1] & 1 == 1 {
-                self.bitmap.push(false);
-            } else {
-                self.bitmap.push(true);
-            }
+            self.bitmap.push(reader[*null_offset - type_size] & 1 != 1);
         }
+        *null_offset -= type_size;
         Ok(())
     }
 

diff --git a/tests/suites/0_stateless/03_dml/03_0018_group_by_nullable.result b/tests/suites/0_stateless/03_dml/03_0018_group_by_nullable.result
@@ -0,0 +1,13 @@
+NULL 2
+0 4
+1 1
+2 3
+NULL 0 2
+0 0 1
+0 1 1
+0 2 1
+0 3 1
+1 NULL 1
+1 0 1
+2 1 1
+2 2 1
diff --git a/tests/suites/0_stateless/03_dml/03_0018_group_by_nullable.sql b/tests/suites/0_stateless/03_dml/03_0018_group_by_nullable.sql
@@ -0,0 +1,5 @@
+CREATE TABLE t(a UInt64, b UInt32) Engine = Fuse;
+INSERT INTO t(a,b)  SELECT if (number % 3 = 1, null, number) as a, number + 3 as b FROM numbers(10);
+SELECT a%3 as c, count(1) as d from t GROUP BY c ORDER BY c, d;
+SELECT a%3 as c, a%4 as d, count(0) as f FROM t GROUP BY c,d ORDER BY c,d,f;
+DROP table t;