[arrow-datafusion] 01/01: Use create_hashes

[dh...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dheres pushed a commit to branch create_hashes_agg
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git

commit c8792d160b3e7e7338781364580622e9530005b2
Author: Daniël Heres <da...@coralogix.com>
AuthorDate: Fri Jul 21 19:11:57 2023 +0200

    Use create_hashes
---
 .../physical_plan/aggregates/group_values/primitive.rs  | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/datafusion/core/src/physical_plan/aggregates/group_values/primitive.rs b/datafusion/core/src/physical_plan/aggregates/group_values/primitive.rs
index 7b8691c67f..3c5ab1ea76 100644
--- a/datafusion/core/src/physical_plan/aggregates/group_values/primitive.rs
+++ b/datafusion/core/src/physical_plan/aggregates/group_values/primitive.rs
@@ -26,6 +26,7 @@ use arrow_schema::DataType;
 use datafusion_common::Result;
 use datafusion_execution::memory_pool::proxy::VecAllocExt;
 use datafusion_physical_expr::EmitTo;
+use datafusion_physical_expr::hash_utils::create_hashes;
 use half::f16;
 use hashbrown::raw::RawTable;
 use std::sync::Arc;
@@ -87,6 +88,8 @@ pub struct GroupValuesPrimitive<T: ArrowPrimitiveType> {
     null_group: Option<usize>,
     /// The values for each group index
     values: Vec<T::Native>,
+    /// buffer to be reused to store hashes
+    hashes_buffer: Vec<u64>,
     /// The random state used to generate hashes
     random_state: RandomState,
 }
@@ -99,6 +102,7 @@ impl<T: ArrowPrimitiveType> GroupValuesPrimitive<T> {
             map: RawTable::with_capacity(128),
             values: Vec::with_capacity(128),
             null_group: None,
+            hashes_buffer: Default::default(),
             random_state: Default::default(),
         }
     }
@@ -112,7 +116,13 @@ where
         assert_eq!(cols.len(), 1);
         groups.clear();
 
-        for v in cols[0].as_primitive::<T>() {
+        // 1.1 Calculate the group keys for the group values
+        let batch_hashes = &mut self.hashes_buffer;
+        batch_hashes.clear();
+        batch_hashes.resize(cols[0].len(), 0);
+        create_hashes(cols, &self.random_state, batch_hashes)?;
+
+        for (v, hash) in cols[0].as_primitive::<T>().into_iter().zip(batch_hashes) {
             let group_id = match v {
                 None => *self.null_group.get_or_insert_with(|| {
                     let group_id = self.values.len();
@@ -121,9 +131,8 @@ where
                 }),
                 Some(key) => {
                     let state = &self.random_state;
-                    let hash = key.hash(state);
                     let insert = self.map.find_or_find_insert_slot(
-                        hash,
+                        *hash,
                         |g| unsafe { self.values.get_unchecked(*g).is_eq(key) },
                         |g| unsafe { self.values.get_unchecked(*g).hash(state) },
                     );
@@ -134,7 +143,7 @@ where
                             Ok(v) => *v.as_ref(),
                             Err(slot) => {
                                 let g = self.values.len();
-                                self.map.insert_in_slot(hash, slot, g);
+                                self.map.insert_in_slot(*hash, slot, g);
                                 self.values.push(key);
                                 g
                             }