implement_group_accumulators_count_distinct_use_hashtable

Author: coderfender

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

@@ -15,40 +15,70 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow::array::{
-    Array, ArrayRef, AsArray, BooleanArray, Int64Array, ListBuilder, PrimitiveBuilder,
-};
-use arrow::datatypes::ArrowPrimitiveType;
+use arrow::array::{ArrayRef, AsArray, BooleanArray, Int64Array, ListArray, PrimitiveArray};
+use arrow::buffer::OffsetBuffer;
+use arrow::datatypes::{ArrowPrimitiveType, Field};
 use datafusion_common::HashSet;
 use datafusion_common::hash_utils::RandomState;
 use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator};
 use std::hash::Hash;
 use std::mem::size_of;
 use std::sync::Arc;
 
+use crate::aggregate::groups_accumulator::accumulate::accumulate;
+
 pub struct PrimitiveDistinctCountGroupsAccumulator<T: ArrowPrimitiveType>
 where
     T::Native: Eq + Hash,
 {
-    /// Count distinct per group.
-    values: Vec<HashSet<T::Native, RandomState>>,
+    seen: HashSet<(usize, T::Native), RandomState>,
+    num_groups: usize,
 }
 
-impl<T: ArrowPrimitiveType> Default for PrimitiveDistinctCountGroupsAccumulator<T>
+impl<T: ArrowPrimitiveType> PrimitiveDistinctCountGroupsAccumulator<T>
 where
     T::Native: Eq + Hash,
 {
-    fn default() -> Self {
-        Self::new()
+    pub fn new() -> Self {
+        Self {
+            seen: HashSet::default(),
+            num_groups: 0,
+        }
+    }
+
+    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> PrimitiveDistinctCountGroupsAccumulator<T>
+impl<T: ArrowPrimitiveType> Default for PrimitiveDistinctCountGroupsAccumulator<T>
 where
     T::Native: Eq + Hash,
 {
-    pub fn new() -> Self {
-        Self { values: Vec::new() }
+    fn default() -> Self {
+        Self::new()
     }
 }
 
@@ -64,47 +94,40 @@ where
         opt_filter: Option<&BooleanArray>,
         total_num_groups: usize,
     ) -> datafusion_common::Result<()> {
-        self.values.resize_with(total_num_groups, HashSet::default);
-        debug_assert_eq!(values.len(), 1, "multiple arguments are not supported");
-
+        debug_assert_eq!(values.len(), 1);
+        self.num_groups = self.num_groups.max(total_num_groups);
         let arr = values[0].as_primitive::<T>();
-
-        for (idx, group_idx) in group_indices.iter().enumerate() {
-            if let Some(filter) = opt_filter
-                && !filter.value(idx)
-            {
-                continue;
-            }
-            if arr.is_valid(idx) {
-                let value = arr.value(idx);
-                self.values[*group_idx].insert(value);
-            }
-        }
-
+        accumulate(group_indices, arr, opt_filter, |group_idx, value| {
+            self.seen.insert((group_idx, value));
+        });
         Ok(())
     }
 
     fn evaluate(&mut self, emit_to: EmitTo) -> datafusion_common::Result<ArrayRef> {
-        let counts: Vec<i64> = emit_to
-            .take_needed(&mut self.values)
-            .iter()
-            .map(|groups| groups.len() as i64)
-            .collect();
-
+        let group_values = self.emit_to_values(emit_to);
+        let counts: Vec<i64> = group_values.iter().map(|v| v.len() as i64).collect();
         Ok(Arc::new(Int64Array::from(counts)))
     }
 
     fn state(&mut self, emit_to: EmitTo) -> datafusion_common::Result<Vec<ArrayRef>> {
-        let hash_sets = emit_to.take_needed(&mut self.values);
-        let mut builder = ListBuilder::new(PrimitiveBuilder::<T>::new());
+        let group_values = self.emit_to_values(emit_to);
 
-        for set in hash_sets {
-            for value in set {
-                builder.values().append_value(value);
-            }
-            builder.append(true);
+        let mut offsets = vec![0i32];
+        let mut all_values = Vec::new();
+        for values in &group_values {
+            all_values.extend(values.iter().copied());
+            offsets.push(all_values.len() as i32);
         }
-        Ok(vec![Arc::new(builder.finish())])
+
+        let values_array = Arc::new(PrimitiveArray::<T>::from_iter_values(all_values));
+        let list_array = ListArray::new(
+            Arc::new(Field::new_list_field(T::DATA_TYPE, true)),
+            OffsetBuffer::new(offsets.into()),
+            values_array,
+            None,
+        );
+
+        Ok(vec![Arc::new(list_array)])
     }
 
     fn merge_batch(
@@ -114,26 +137,23 @@ where
         _opt_filter: Option<&BooleanArray>,
         total_num_groups: usize,
     ) -> datafusion_common::Result<()> {
-        self.values.resize_with(total_num_groups, HashSet::default);
+        debug_assert_eq!(values.len(), 1);
+        self.num_groups = self.num_groups.max(total_num_groups);
         let list_array = values[0].as_list::<i32>();
 
         for (row_idx, group_idx) in group_indices.iter().enumerate() {
             let inner = list_array.value(row_idx);
-            let inner_set = inner.as_primitive::<T>();
-            for i in 0..inner.len() {
-                self.values[*group_idx].insert(inner_set.value(i));
+            let inner_arr = inner.as_primitive::<T>();
+            for value in inner_arr.values().iter() {
+                self.seen.insert((*group_idx, *value));
             }
         }
+
         Ok(())
     }
 
     fn size(&self) -> usize {
         size_of::<Self>()
-            + self.values.capacity() * size_of::<HashSet<T::Native, RandomState>>()
-            + self
-                .values
-                .iter()
-                .map(|s| s.capacity() * size_of::<T::Native>())
-                .sum::<usize>()
+            + self.seen.capacity() * (size_of::<(usize, T::Native)>() + size_of::<u64>())
     }
 }