fix batch equal to logic.

Author: Rachelint

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

@@ -28,29 +28,37 @@ struct AggregationHashTable<T: AggregationHashTableEntry> {
     /// Pending indices needed to batch compare after
     /// If equal after comparing, this group exist
     /// If not equal after comparing, prob the next bucket
-    pending_equal_to_indices: Vec<usize>,
+    pending_equal_to_ctxs: Vec<EqualToContext>,
 
     /// The non equal indices in one round
     non_equal_indices: Vec<usize>,
 }
 
+pub struct EqualToContext {
+    row_index: usize,
+    group_index: usize,
+}
+
 #[derive(Debug, Clone)]
 pub enum AggregationHashTableEntryState<T: AggregationHashTableEntry> {
     Entry(T),
-    PlaceHolder,
+    PlaceHolder(usize),
     Empty,
 }
 
 trait AggregationHashTableEntry: fmt::Debug + Clone + Default {
     fn check_hash(&self, target_hash: u64) -> bool;
 
-    fn get_hash(&self, random_state: &RandomState) -> u64;
+    fn hash_value(&self, random_state: &RandomState) -> u64;
+
+    fn group_index(&self) -> usize;
 }
 
 impl<T: AggregationHashTableEntry> AggregationHashTable<T> {
     fn get_or_create_groups<A, E, C>(
         &mut self,
         cols: &[ArrayRef],
+        mut current_total_groups: usize,
         batch_hashes: &[u64],
         random_state: &RandomState,
         mut batch_append: A,
@@ -59,7 +67,7 @@ impl<T: AggregationHashTableEntry> AggregationHashTable<T> {
         groups: &mut Vec<usize>,
     ) where
         A: FnMut(&[ArrayRef], &[usize], &mut Vec<usize>),
-        E: FnMut(&[ArrayRef], &[usize], &mut Vec<usize>, &mut Vec<usize>),
+        E: FnMut(&[ArrayRef], &[EqualToContext], &mut Vec<usize>, &mut Vec<usize>),
         C: Fn(usize, u64) -> T,
     {
         let num_input_rows = cols[0].len();
@@ -78,7 +86,7 @@ impl<T: AggregationHashTableEntry> AggregationHashTable<T> {
         self.remaining_indices.extend(0..num_input_rows);
         while self.remaining_indices.len() > 0 {
             self.pending_append_indices.clear();
-            self.pending_equal_to_indices.clear();
+            self.pending_equal_to_ctxs.clear();
             self.non_equal_indices.clear();
 
             // 3. Find entry in `raw_table` first, and check:
@@ -92,19 +100,31 @@ impl<T: AggregationHashTableEntry> AggregationHashTable<T> {
 
                 let table_entry = match &table_entry {
                     AggregationHashTableEntryState::Empty => {
-                        *table_entry = AggregationHashTableEntryState::PlaceHolder;
+                        self.pending_append_indices.push(row_idx);
+                        *table_entry = AggregationHashTableEntryState::PlaceHolder(
+                            current_total_groups,
+                        );
+                        current_total_groups += 1;
                         continue;
                     }
-                    AggregationHashTableEntryState::PlaceHolder => {
-                        self.non_equal_indices.push(row_idx);
+                    AggregationHashTableEntryState::PlaceHolder(group_index) => {
+                        let equal_to_ctx = EqualToContext {
+                            row_index: row_idx,
+                            group_index: *group_index,
+                        };
+                        self.pending_equal_to_ctxs.push(equal_to_ctx);
                         continue;
                     }
                     AggregationHashTableEntryState::Entry(entry) => entry,
                 };
 
                 let target_hash = batch_hashes[row_idx];
                 if table_entry.check_hash(target_hash) {
-                    self.pending_equal_to_indices.push(row_idx);
+                    let equal_to_ctx = EqualToContext {
+                        row_index: row_idx,
+                        group_index: table_entry.group_index(),
+                    };
+                    self.pending_equal_to_ctxs.push(equal_to_ctx);
                 } else {
                     self.non_equal_indices.push(row_idx);
                 }
@@ -122,7 +142,7 @@ impl<T: AggregationHashTableEntry> AggregationHashTable<T> {
                     let offset = self.table_offsets[row_idx];
                     debug_assert!(matches!(
                         &self.raw_table[offset],
-                        AggregationHashTableEntryState::PlaceHolder
+                        AggregationHashTableEntryState::PlaceHolder(_)
                     ));
                     self.raw_table[offset] =
                         AggregationHashTableEntryState::Entry(new_entry(group_idx, hash));
@@ -133,10 +153,10 @@ impl<T: AggregationHashTableEntry> AggregationHashTable<T> {
             }
 
             // 5. Batch equal to
-            if !self.pending_equal_to_indices.is_empty() {
+            if !self.pending_equal_to_ctxs.is_empty() {
                 batch_equal_to(
                     &cols,
-                    &self.pending_equal_to_indices,
+                    &self.pending_equal_to_ctxs,
                     &mut self.non_equal_indices,
                     groups,
                 );
@@ -179,10 +199,10 @@ impl<T: AggregationHashTableEntry> AggregationHashTable<T> {
         for state in old_raw_table {
             debug_assert!(!matches!(
                 &state,
-                AggregationHashTableEntryState::PlaceHolder
+                AggregationHashTableEntryState::PlaceHolder(_)
             ));
             if let AggregationHashTableEntryState::Entry(entry) = &state {
-                let hash = entry.get_hash(random_state);
+                let hash = entry.hash_value(random_state);
                 let offset = (hash & self.bit_mask) as usize;
                 self.raw_table[offset] = state;
             }