Implement GroupColumn support for StringView / ByteView (faster…

… grouping performance) (#12809) * define `ByteGroupValueViewBuilder`. * impl append. * impl equal to. * fix compile. * fix comments. * impl take_n. * impl build. * impl rest functions in `GroupColumn`. * fix output when panic. * add e2e sql tests. * add unit tests. * switch to a really elegant style codes from alamb. * fix take_n. * improve comments. * fix compile. * fix clippy. * define more testcases in `test_byte_view_take_n`. * connect up. * fix doc. * Do not re-validate output is utf8 * switch to unchecked when building array. * improve naming. * use let else to make the codes clearer. * fix typo. * improve unit test coverage for `ByteViewGroupValueBuilder`. --------- Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
apache · Oct 16, 2024 · c311cf5 · c311cf5
1 parent e5cdc17
commit c311cf5
Show file tree

Hide file tree

Showing 3 changed files with 755 additions and 6 deletions.
diff --git a/datafusion/physical-plan/src/aggregates/group_values/column.rs b/datafusion/physical-plan/src/aggregates/group_values/column.rs
@@ -16,14 +16,16 @@
 // under the License.
 
 use crate::aggregates::group_values::group_column::{
-    ByteGroupValueBuilder, GroupColumn, PrimitiveGroupValueBuilder,
+    ByteGroupValueBuilder, ByteViewGroupValueBuilder, GroupColumn,
+    PrimitiveGroupValueBuilder,
 };
 use crate::aggregates::group_values::GroupValues;
 use ahash::RandomState;
 use arrow::compute::cast;
 use arrow::datatypes::{
-    Date32Type, Date64Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type,
-    Int8Type, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
+    BinaryViewType, Date32Type, Date64Type, Float32Type, Float64Type, Int16Type,
+    Int32Type, Int64Type, Int8Type, StringViewType, UInt16Type, UInt32Type, UInt64Type,
+    UInt8Type,
 };
 use arrow::record_batch::RecordBatch;
 use arrow_array::{Array, ArrayRef};
@@ -119,6 +121,8 @@ impl GroupValuesColumn {
                 | DataType::LargeBinary
                 | DataType::Date32
                 | DataType::Date64
+                | DataType::Utf8View
+                | DataType::BinaryView
         )
     }
 }
@@ -184,6 +188,14 @@ impl GroupValues for GroupValuesColumn {
                         let b = ByteGroupValueBuilder::<i64>::new(OutputType::Binary);
                         v.push(Box::new(b) as _)
                     }
+                    &DataType::Utf8View => {
+                        let b = ByteViewGroupValueBuilder::<StringViewType>::new();
+                        v.push(Box::new(b) as _)
+                    }
+                    &DataType::BinaryView => {
+                        let b = ByteViewGroupValueBuilder::<BinaryViewType>::new();
+                        v.push(Box::new(b) as _)
+                    }
                     dt => {
                         return not_impl_err!("{dt} not supported in GroupValuesColumn")
                     }