implement_group_accumulators_count_distinct_use_hashtable

Author: coderfender

datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs

@@ -48,31 +48,6 @@ where
         }
     }
 
-    fn emit_to_values(&mut self, emit_to: EmitTo) -> Vec<Vec<T::Native>> {
-        let num_emitted = match emit_to {
-            EmitTo::All => self.num_groups,
-            EmitTo::First(n) => n,
-        };
-
-        let mut group_values: Vec<Vec<T::Native>> = vec![Vec::new(); num_emitted];
-        let mut remaining = HashSet::default();
-
-        for (group_idx, value) in self.seen.drain() {
-            if group_idx < num_emitted {
-                group_values[group_idx].push(value);
-            } else {
-                remaining.insert((group_idx - num_emitted, value));
-            }
-        }
-
-        self.seen = remaining;
-        match emit_to {
-            EmitTo::All => self.num_groups = 0,
-            EmitTo::First(n) => self.num_groups = self.num_groups.saturating_sub(n),
-        }
-
-        group_values
-    }
 }
 
 impl<T: ArrowPrimitiveType> Default for PrimitiveDistinctCountGroupsAccumulator<T>
@@ -106,13 +81,60 @@ where
     }
 
     fn evaluate(&mut self, emit_to: EmitTo) -> datafusion_common::Result<ArrayRef> {
-        let group_values = self.emit_to_values(emit_to);
-        let counts: Vec<i64> = group_values.iter().map(|v| v.len() as i64).collect();
+        let num_emitted = match emit_to {
+            EmitTo::All => self.num_groups,
+            EmitTo::First(n) => n,
+        };
+
+        let mut counts = vec![0i64; num_emitted];
+
+        if matches!(emit_to, EmitTo::All) {
+            for &(group_idx, _) in self.seen.iter() {
+                counts[group_idx] += 1;
+            }
+            self.seen.clear();
+            self.num_groups = 0;
+        } else {
+            let mut remaining = HashSet::default();
+            for (group_idx, value) in self.seen.drain() {
+                if group_idx < num_emitted {
+                    counts[group_idx] += 1;
+                } else {
+                    remaining.insert((group_idx - num_emitted, value));
+                }
+            }
+            self.seen = remaining;
+            self.num_groups = self.num_groups.saturating_sub(num_emitted);
+        }
+
         Ok(Arc::new(Int64Array::from(counts)))
     }
 
     fn state(&mut self, emit_to: EmitTo) -> datafusion_common::Result<Vec<ArrayRef>> {
-        let group_values = self.emit_to_values(emit_to);
+        let num_emitted = match emit_to {
+            EmitTo::All => self.num_groups,
+            EmitTo::First(n) => n,
+        };
+
+        let mut group_values: Vec<Vec<T::Native>> = vec![Vec::new(); num_emitted];
+
+        if matches!(emit_to, EmitTo::All) {
+            for (group_idx, value) in self.seen.drain() {
+                group_values[group_idx].push(value);
+            }
+            self.num_groups = 0;
+        } else {
+            let mut remaining = HashSet::default();
+            for (group_idx, value) in self.seen.drain() {
+                if group_idx < num_emitted {
+                    group_values[group_idx].push(value);
+                } else {
+                    remaining.insert((group_idx - num_emitted, value));
+                }
+            }
+            self.seen = remaining;
+            self.num_groups = self.num_groups.saturating_sub(num_emitted);
+        }
 
         let mut offsets = vec![0i32];
         let mut all_values = Vec::new();