implement to Arrow's builder

Signed-off-by: jayzhan211 <[email protected]>

Author: jayzhan211

benchmarks/queries/clickbench/queries.sql

@@ -35,4 +35,9 @@ SELECT "URL", COUNT(*) AS c FROM hits GROUP BY "URL" ORDER BY c DESC LIMIT 10;
 SELECT 1, "URL", COUNT(*) AS c FROM hits GROUP BY 1, "URL" ORDER BY c DESC LIMIT 10;
 SELECT "ClientIP", "ClientIP" - 1, "ClientIP" - 2, "ClientIP" - 3, COUNT(*) AS c FROM hits GROUP BY "ClientIP", "ClientIP" - 1, "ClientIP" - 2, "ClientIP" - 3 ORDER BY c DESC LIMIT 10;
 SELECT "URL", COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate"::INT::DATE >= '2013-07-01' AND "EventDate"::INT::DATE <= '2013-07-31' AND "DontCountHits" = 0 AND "IsRefresh" = 0 AND "URL" <> '' GROUP BY "URL" ORDER BY PageViews DESC LIMIT 10;
-SELECT "Title", COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate"::INT::DATE >= '2013-07-01' AND "EventDate"::INT::DATE <= '2013-07-31' AND "DontCountHits" = 0 AND "IsRefresh" = 0 AND "Title" <> '' GROUP BY "Title" ORDER BY PageViews DESC LIMIT 10;
+SELECT "Title", COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate"::INT::DATE >= '2013-07-01' AND "EventDate"::INT::DATE <= '2013-07-31' AND "DontCountHits" = 0 AND "IsRefresh" = 0 AND "Title" <> '' GROUP BY "Title" ORDER BY PageViews DESC LIMIT 10;
+SELECT "URL", COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate"::INT::DATE >= '2013-07-01' AND "EventDate"::INT::DATE <= '2013-07-31' AND "IsRefresh" = 0 AND "IsLink" <> 0 AND "IsDownload" = 0 GROUP BY "URL" ORDER BY PageViews DESC LIMIT 10 OFFSET 1000;
+SELECT "TraficSourceID", "SearchEngineID", "AdvEngineID", CASE WHEN ("SearchEngineID" = 0 AND "AdvEngineID" = 0) THEN "Referer" ELSE '' END AS Src, "URL" AS Dst, COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate"::INT::DATE >= '2013-07-01' AND "EventDate"::INT::DATE <= '2013-07-31' AND "IsRefresh" = 0 GROUP BY "TraficSourceID", "SearchEngineID", "AdvEngineID", Src, Dst ORDER BY PageViews DESC LIMIT 10 OFFSET 1000;
+SELECT "URLHash", "EventDate"::INT::DATE, COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate"::INT::DATE >= '2013-07-01' AND "EventDate"::INT::DATE <= '2013-07-31' AND "IsRefresh" = 0 AND "TraficSourceID" IN (-1, 6) AND "RefererHash" = 3594120000172545465 GROUP BY "URLHash", "EventDate"::INT::DATE ORDER BY PageViews DESC LIMIT 10 OFFSET 100;
+SELECT "WindowClientWidth", "WindowClientHeight", COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate"::INT::DATE >= '2013-07-01' AND "EventDate"::INT::DATE <= '2013-07-31' AND "IsRefresh" = 0 AND "DontCountHits" = 0 AND "URLHash" = 2868770270353813622 GROUP BY "WindowClientWidth", "WindowClientHeight" ORDER BY PageViews DESC LIMIT 10 OFFSET 10000;
+SELECT DATE_TRUNC('minute', to_timestamp_seconds("EventTime")) AS M, COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate"::INT::DATE >= '2013-07-14' AND "EventDate"::INT::DATE <= '2013-07-15' AND "IsRefresh" = 0 AND "DontCountHits" = 0 GROUP BY DATE_TRUNC('minute', to_timestamp_seconds("EventTime")) ORDER BY DATE_TRUNC('minute', M) LIMIT 10 OFFSET 1000;

datafusion/physical-expr-common/src/group_value_row.rs

@@ -18,8 +18,11 @@
 use arrow::array::BooleanBufferBuilder;
 use arrow::array::BufferBuilder;
 use arrow::array::GenericBinaryArray;
+use arrow::array::GenericBinaryBuilder;
 use arrow::array::GenericStringArray;
+use arrow::array::GenericStringBuilder;
 use arrow::array::OffsetSizeTrait;
+use arrow::array::PrimitiveBuilder;
 use arrow::buffer::NullBuffer;
 use arrow::buffer::OffsetBuffer;
 use arrow::buffer::ScalarBuffer;
@@ -37,13 +40,61 @@ use crate::binary_map::INITIAL_BUFFER_CAPACITY;
 
 use std::sync::Arc;
 
+pub trait ArrayEqV2: Send + Sync {
+    fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool;
+    fn append_val(&mut self, array: &ArrayRef, row: usize);
+    fn len(&self) -> usize;
+    fn build(&mut self) -> ArrayRef;
+}
+
 pub trait ArrayEq: Send + Sync {
     fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool;
     fn append_val(&mut self, array: &ArrayRef, row: usize);
     fn len(&self) -> usize;
     fn build(self: Box<Self>) -> ArrayRef;
 }
 
+impl<T> ArrayEqV2 for PrimitiveBuilder<T>
+where
+    T: ArrowPrimitiveType,
+{
+    fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool {
+        let arr = array.as_primitive::<T>();
+
+        if let Some(nulls) = self.validity_slice() {
+            let null_slice_index = lhs_row / 8;
+            let null_bit_map_index = lhs_row % 8;
+            let is_elem_null = ((nulls[null_slice_index] >> null_bit_map_index) & 1) == 1;
+            if is_elem_null {
+                return arr.is_null(rhs_row);
+            } else if arr.is_null(rhs_row) {
+                return false;
+            }
+        }
+
+        let elem = self.values_slice()[lhs_row];
+        elem == arr.value(rhs_row)
+    }
+
+    fn append_val(&mut self, array: &ArrayRef, row: usize) {
+        let arr = array.as_primitive::<T>();
+        if arr.is_null(row) {
+            self.append_null();
+        } else {
+            let elem = arr.value(row);
+            self.append_value(elem);
+        }
+    }
+
+    fn len(&self) -> usize {
+        self.values_slice().len()
+    }
+
+    fn build(&mut self) -> ArrayRef {
+        Arc::new(self.finish())
+    }
+}
+
 pub struct PrimitiveGroupValueBuilder<T: ArrowPrimitiveType>(Vec<Option<T::Native>>);
 
 impl<T: ArrowPrimitiveType> PrimitiveGroupValueBuilder<T> {
@@ -127,6 +178,99 @@ impl<T: ArrowPrimitiveType> ArrayEq for PrimitiveGroupValueBuilder<T> {
 //     }
 // }
 
+impl<O> ArrayEqV2 for GenericStringBuilder<O>
+where
+    O: OffsetSizeTrait,
+{
+    fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool {
+        let arr = array.as_bytes::<GenericStringType<O>>();
+        if let Some(nulls) = self.validity_slice() {
+            let null_slice_index = lhs_row / 8;
+            let null_bit_map_index = lhs_row % 8;
+
+            let is_lhs_null = ((nulls[null_slice_index] >> null_bit_map_index) & 1) == 1;
+            if is_lhs_null {
+                return arr.is_null(rhs_row);
+            } else if arr.is_null(rhs_row) {
+                return false;
+            }
+        }
+
+        let rhs_elem: &[u8] = arr.value(rhs_row).as_ref();
+        let rhs_elem_len = arr.value_length(rhs_row).as_usize();
+        assert_eq!(rhs_elem_len, rhs_elem.len());
+        let l = O::as_usize(self.offsets_slice()[lhs_row]);
+        let r = O::as_usize(self.offsets_slice()[lhs_row + 1]);
+        let existing_elem = &self.values_slice()[l..r];
+        existing_elem.len() == rhs_elem.len() && rhs_elem == existing_elem
+    }
+
+    fn append_val(&mut self, array: &ArrayRef, row: usize) {
+        let arr = array.as_string::<O>();
+        if arr.is_null(row) {
+            self.append_null();
+            return;
+        }
+
+        let value = arr.value(row);
+        self.append_value(value);
+    }
+
+    fn len(&self) -> usize {
+        self.offsets_slice().len() - 1
+    }
+
+    fn build(&mut self) -> ArrayRef {
+        Arc::new(self.finish())
+    }
+}
+
+impl<O> ArrayEqV2 for GenericBinaryBuilder<O>
+where
+    O: OffsetSizeTrait,
+{
+    fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool {
+        let arr = array.as_bytes::<GenericBinaryType<O>>();
+        if let Some(nulls) = self.validity_slice() {
+            let null_slice_index = lhs_row / 8;
+            let null_bit_map_index = lhs_row % 8;
+            let is_lhs_null = ((nulls[null_slice_index] >> null_bit_map_index) & 1) == 1;
+            if is_lhs_null {
+                return arr.is_null(rhs_row);
+            } else if arr.is_null(rhs_row) {
+                return false;
+            }
+        }
+
+        let rhs_elem: &[u8] = arr.value(rhs_row).as_ref();
+        let rhs_elem_len = arr.value_length(rhs_row).as_usize();
+        assert_eq!(rhs_elem_len, rhs_elem.len());
+        let l = O::as_usize(self.offsets_slice()[lhs_row]);
+        let r = O::as_usize(self.offsets_slice()[lhs_row + 1]);
+        let existing_elem = &self.values_slice()[l..r];
+        existing_elem.len() == rhs_elem.len() && rhs_elem == existing_elem
+    }
+
+    fn append_val(&mut self, array: &ArrayRef, row: usize) {
+        let arr = array.as_binary::<O>();
+        if arr.is_null(row) {
+            self.append_null();
+            return;
+        }
+
+        let value: &[u8] = arr.value(row).as_ref();
+        self.append_value(value);
+    }
+
+    fn len(&self) -> usize {
+        self.values_slice().len()
+    }
+
+    fn build(&mut self) -> ArrayRef {
+        Arc::new(self.finish())
+    }
+}
+
 pub struct ByteGroupValueBuilderNaive<O>
 where
     O: OffsetSizeTrait,
@@ -622,3 +766,26 @@ where
 //         self.offset_or_inline..self.offset_or_inline + self.len.as_usize()
 //     }
 // }
+
+#[cfg(test)]
+mod tests {
+    use arrow::{array::GenericByteBuilder, datatypes::GenericStringType};
+
+    #[test]
+    fn test123() {
+        let mut a = GenericByteBuilder::<GenericStringType<i32>>::new();
+        a.append_null();
+        a.append_value("a");
+        a.append_null();
+        a.append_value("bc");
+        a.append_value("def");
+        a.append_null();
+
+        let s = a.validity_slice();
+        println!("s: {:?}", s);
+        let v = a.values_slice();
+        let o = a.offsets_slice();
+        println!("v: {:?}", v);
+        println!("o: {:?}", o);
+    }
+}

datafusion/physical-plan/src/aggregates/group_values/mod.rs

@@ -117,7 +117,9 @@ fn has_row_like_feature(data_type: &DataType) -> bool {
         | DataType::UInt32
         | DataType::UInt64
         | DataType::Utf8
-        | DataType::LargeUtf8 => true,
+        | DataType::LargeUtf8
+        | DataType::Date32
+        | DataType::Date64 => true,
         _ => false,
     }
 }

datafusion/physical-plan/src/aggregates/group_values/row_like.rs

@@ -15,12 +15,15 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::ops::DerefMut;
+
 use crate::aggregates::group_values::GroupValues;
 use ahash::RandomState;
+use arrow::array::{GenericStringBuilder, PrimitiveBuilder};
 use arrow::compute::cast;
 use arrow::datatypes::{
-    Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, UInt16Type,
-    UInt32Type, UInt64Type, UInt8Type,
+    Date32Type, Date64Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type,
+    Int8Type, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
 };
 use arrow::record_batch::RecordBatch;
 use arrow::row::{RowConverter, Rows, SortField};
@@ -30,12 +33,11 @@ use datafusion_common::hash_utils::create_hashes;
 use datafusion_common::{internal_err, DataFusionError, Result};
 use datafusion_execution::memory_pool::proxy::{RawTableAllocExt, VecAllocExt};
 use datafusion_expr::EmitTo;
-use datafusion_physical_expr::binary_map::OutputType;
-use datafusion_physical_expr_common::group_value_row::{
-    ArrayEq, ByteGroupValueBuilderNaive, PrimitiveGroupValueBuilder,
-};
+use datafusion_physical_expr_common::group_value_row::ArrayEqV2;
 use hashbrown::raw::RawTable;
 
+pub(super) const INITIAL_CAPACITY: usize = 8 * 1024;
+
 /// A [`GroupValues`] making use of [`Rows`]
 pub struct GroupValuesRowLike {
     /// The output schema
@@ -75,7 +77,7 @@ pub struct GroupValuesRowLike {
 
     /// Random state for creating hashes
     random_state: RandomState,
-    group_values_v2: Option<Vec<Box<dyn ArrayEq>>>,
+    group_values_v2: Option<Vec<Box<dyn ArrayEqV2>>>,
 }
 
 impl GroupValuesRowLike {
@@ -131,53 +133,65 @@ impl GroupValues for GroupValuesRowLike {
                 for (i, f) in self.schema.fields().iter().enumerate() {
                     match f.data_type() {
                         &DataType::Int8 => {
-                            let b = PrimitiveGroupValueBuilder::<Int8Type>::new();
+                            let b = PrimitiveBuilder::<Int8Type>::new();
                             v.push(Box::new(b) as _)
                         }
                         &DataType::Int16 => {
-                            let b = PrimitiveGroupValueBuilder::<Int16Type>::new();
+                            let b = PrimitiveBuilder::<Int16Type>::new();
                             v.push(Box::new(b) as _)
                         }
                         &DataType::Int32 => {
-                            let b = PrimitiveGroupValueBuilder::<Int32Type>::new();
+                            let b = PrimitiveBuilder::<Int32Type>::new();
                             v.push(Box::new(b) as _)
                         }
                         &DataType::Int64 => {
-                            let b = PrimitiveGroupValueBuilder::<Int64Type>::new();
+                            let b = PrimitiveBuilder::<Int64Type>::new();
                             v.push(Box::new(b) as _)
                         }
                         &DataType::UInt8 => {
-                            let b = PrimitiveGroupValueBuilder::<UInt8Type>::new();
+                            let b = PrimitiveBuilder::<UInt8Type>::new();
                             v.push(Box::new(b) as _)
                         }
                         &DataType::UInt16 => {
-                            let b = PrimitiveGroupValueBuilder::<UInt16Type>::new();
+                            let b = PrimitiveBuilder::<UInt16Type>::new();
                             v.push(Box::new(b) as _)
                         }
                         &DataType::UInt32 => {
-                            let b = PrimitiveGroupValueBuilder::<UInt32Type>::new();
+                            let b = PrimitiveBuilder::<UInt32Type>::new();
                             v.push(Box::new(b) as _)
                         }
                         &DataType::UInt64 => {
-                            let b = PrimitiveGroupValueBuilder::<UInt64Type>::new();
+                            let b = PrimitiveBuilder::<UInt64Type>::new();
                             v.push(Box::new(b) as _)
                         }
                         &DataType::Float32 => {
-                            let b = PrimitiveGroupValueBuilder::<Float32Type>::new();
+                            let b = PrimitiveBuilder::<Float32Type>::new();
                             v.push(Box::new(b) as _)
                         }
                         &DataType::Float64 => {
-                            let b = PrimitiveGroupValueBuilder::<Float64Type>::new();
+                            let b = PrimitiveBuilder::<Float64Type>::new();
+                            v.push(Box::new(b) as _)
+                        }
+                        &DataType::Date32 => {
+                            let b = PrimitiveBuilder::<Date32Type>::new();
+                            v.push(Box::new(b) as _)
+                        }
+                        &DataType::Date64 => {
+                            let b = PrimitiveBuilder::<Date64Type>::new();
                             v.push(Box::new(b) as _)
                         }
                         &DataType::Utf8 => {
-                            let b =
-                                ByteGroupValueBuilderNaive::<i32>::new(OutputType::Utf8);
+                            let b = GenericStringBuilder::<i32>::with_capacity(
+                                INITIAL_CAPACITY,
+                                INITIAL_CAPACITY,
+                            );
                             v.push(Box::new(b) as _)
                         }
                         &DataType::LargeUtf8 => {
-                            let b =
-                                ByteGroupValueBuilderNaive::<i64>::new(OutputType::Utf8);
+                            let b = GenericStringBuilder::<i64>::with_capacity(
+                                INITIAL_CAPACITY,
+                                INITIAL_CAPACITY,
+                            );
                             v.push(Box::new(b) as _)
                         }
                         dt => todo!("{dt} not impl"),
@@ -211,7 +225,7 @@ impl GroupValues for GroupValuesRowLike {
                 // && group_rows.row(row) == group_values.row(*group_idx)
 
                 fn compare_equal(
-                    arry_eq: &dyn ArrayEq,
+                    arry_eq: &dyn ArrayEqV2,
                     lhs_row: usize,
                     array: &ArrayRef,
                     rhs_row: usize,
@@ -311,7 +325,10 @@ impl GroupValues for GroupValuesRowLike {
             EmitTo::All => {
                 let output = group_values_v2
                     .into_iter()
-                    .map(|v| v.build())
+                    .map(|mut v| {
+                        let p = v.deref_mut().build();
+                        p
+                    })
                     .collect::<Vec<_>>();
                 // let output = self.row_converter.convert_rows(&group_values)?;
                 // group_values.clear();
@@ -324,8 +341,15 @@ impl GroupValues for GroupValuesRowLike {
 
                 // println!("to first n");
                 let len = group_values_v2.len();
-                let first_n: Vec<Box<dyn ArrayEq>> = group_values_v2.drain(..n).collect();
-                let output = first_n.into_iter().map(|v| v.build()).collect::<Vec<_>>();
+                let first_n: Vec<Box<dyn ArrayEqV2>> =
+                    group_values_v2.drain(..n).collect();
+                let output = first_n
+                    .into_iter()
+                    .map(|mut v| {
+                        let p = v.deref_mut().build();
+                        p
+                    })
+                    .collect::<Vec<_>>();
                 assert_eq!(len, group_values_v2.len() + n);
                 self.group_values_v2 = Some(group_values_v2);