expl_tuple_vec_on_demand_sort

expl_tuple_vec_on_demand_sort

Author: B Vadlamani

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

@@ -20,38 +20,91 @@ use arrow::array::{
 };
 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;
 
+/// Trait for packing (group_idx, value) into a single sortable integer
+pub trait Packable: Copy + Send {
+    type Packed: Ord + Copy + Send;
+    fn pack(group_idx: usize, value: Self) -> Self::Packed;
+    fn unpack(packed: Self::Packed) -> (usize, Self);
+}
+
+macro_rules! impl_packable_signed {
+    ($native:ty, $unsigned:ty, $packed:ty, $bits:expr) => {
+        impl Packable for $native {
+            type Packed = $packed;
+            #[inline]
+            fn pack(group_idx: usize, value: Self) -> $packed {
+                let val = (value as $unsigned ^ (1 << ($bits - 1))) as $packed;
+                ((group_idx as $packed) << $bits) | val
+            }
+            #[inline]
+            fn unpack(packed: $packed) -> (usize, Self) {
+                let group = (packed >> $bits) as usize;
+                let val = ((packed as $unsigned) ^ (1 << ($bits - 1))) as $native;
+                (group, val)
+            }
+        }
+    };
+}
+
+macro_rules! impl_packable_unsigned {
+    ($native:ty, $packed:ty, $bits:expr) => {
+        impl Packable for $native {
+            type Packed = $packed;
+            #[inline]
+            fn pack(group_idx: usize, value: Self) -> $packed {
+                ((group_idx as $packed) << $bits) | (value as $packed)
+            }
+            #[inline]
+            fn unpack(packed: $packed) -> (usize, Self) {
+                ((packed >> $bits) as usize, packed as $native)
+            }
+        }
+    };
+}
+
+impl_packable_signed!(i64, u64, u128, 64);
+impl_packable_signed!(i32, u32, u64, 32);
+impl_packable_signed!(i16, u16, u64, 16);
+impl_packable_signed!(i8, u8, u64, 8);
+
+impl_packable_unsigned!(u64, u128, 64);
+impl_packable_unsigned!(u32, u64, 32);
+impl_packable_unsigned!(u16, u64, 16);
+impl_packable_unsigned!(u8, u64, 8);
+
+/// A `GroupsAccumulator` for COUNT(DISTINCT) on primitive types.
+///
+/// Uses a flat buffer with packed integers: (group_idx, value) packed into
+/// a single sortable integer. Sort and dedup at evaluate time.
 pub struct PrimitiveDistinctCountGroupsAccumulator<T: ArrowPrimitiveType>
 where
-    T::Native: Eq + Hash,
+    T::Native: Packable,
 {
-    seen: HashSet<(usize, T::Native), RandomState>,
+    buffer: Vec<<T::Native as Packable>::Packed>,
     num_groups: usize,
 }
 
 impl<T: ArrowPrimitiveType> PrimitiveDistinctCountGroupsAccumulator<T>
 where
-    T::Native: Eq + Hash,
+    T::Native: Packable,
 {
     pub fn new() -> Self {
         Self {
-            seen: HashSet::default(),
+            buffer: Vec::new(),
             num_groups: 0,
         }
     }
 }
 
 impl<T: ArrowPrimitiveType> Default for PrimitiveDistinctCountGroupsAccumulator<T>
 where
-    T::Native: Eq + Hash,
+    T::Native: Packable,
 {
     fn default() -> Self {
         Self::new()
@@ -61,7 +114,7 @@ where
 impl<T: ArrowPrimitiveType + Send + std::fmt::Debug> GroupsAccumulator
     for PrimitiveDistinctCountGroupsAccumulator<T>
 where
-    T::Native: Eq + Hash,
+    T::Native: Packable,
 {
     fn update_batch(
         &mut self,
@@ -73,8 +126,11 @@ where
         debug_assert_eq!(values.len(), 1);
         self.num_groups = self.num_groups.max(total_num_groups);
         let arr = values[0].as_primitive::<T>();
+
+        self.buffer.reserve(arr.len());
+
         accumulate(group_indices, arr, opt_filter, |group_idx, value| {
-            self.seen.insert((group_idx, value));
+            self.buffer.push(T::Native::pack(group_idx, value));
         });
         Ok(())
     }
@@ -85,24 +141,31 @@ where
             EmitTo::First(n) => n,
         };
 
+        self.buffer.sort_unstable();
+
         let mut counts = vec![0i64; num_emitted];
 
         if matches!(emit_to, EmitTo::All) {
-            for &(group_idx, _) in self.seen.iter() {
+            self.buffer.dedup();
+            for &packed in &self.buffer {
+                let (group_idx, _) = T::Native::unpack(packed);
                 counts[group_idx] += 1;
             }
-            self.seen.clear();
+            self.buffer.clear();
             self.num_groups = 0;
         } else {
-            let mut remaining = HashSet::default();
-            for (group_idx, value) in self.seen.drain() {
+            self.buffer.dedup();
+            let mut remaining = Vec::new();
+
+            for &packed in &self.buffer {
+                let (group_idx, value) = T::Native::unpack(packed);
                 if group_idx < num_emitted {
                     counts[group_idx] += 1;
                 } else {
-                    remaining.insert((group_idx - num_emitted, value));
+                    remaining.push(T::Native::pack(group_idx - num_emitted, value));
                 }
             }
-            self.seen = remaining;
+            self.buffer = remaining;
             self.num_groups = self.num_groups.saturating_sub(num_emitted);
         }
 
@@ -115,23 +178,29 @@ where
             EmitTo::First(n) => n,
         };
 
+        self.buffer.sort_unstable();
+        self.buffer.dedup();
+
         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() {
+            for &packed in &self.buffer {
+                let (group_idx, value) = T::Native::unpack(packed);
                 group_values[group_idx].push(value);
             }
+            self.buffer.clear();
             self.num_groups = 0;
         } else {
-            let mut remaining = HashSet::default();
-            for (group_idx, value) in self.seen.drain() {
+            let mut remaining = Vec::new();
+            for &packed in &self.buffer {
+                let (group_idx, value) = T::Native::unpack(packed);
                 if group_idx < num_emitted {
                     group_values[group_idx].push(value);
                 } else {
-                    remaining.insert((group_idx - num_emitted, value));
+                    remaining.push(T::Native::pack(group_idx - num_emitted, value));
                 }
             }
-            self.seen = remaining;
+            self.buffer = remaining;
             self.num_groups = self.num_groups.saturating_sub(num_emitted);
         }
 
@@ -167,8 +236,9 @@ where
         for (row_idx, group_idx) in group_indices.iter().enumerate() {
             let inner = list_array.value(row_idx);
             let inner_arr = inner.as_primitive::<T>();
+            self.buffer.reserve(inner_arr.len());
             for value in inner_arr.values().iter() {
-                self.seen.insert((*group_idx, *value));
+                self.buffer.push(T::Native::pack(*group_idx, *value));
             }
         }
 
@@ -177,6 +247,6 @@ where
 
     fn size(&self) -> usize {
         size_of::<Self>()
-            + self.seen.capacity() * (size_of::<(usize, T::Native)>() + size_of::<u64>())
+            + self.buffer.capacity() * size_of::<<T::Native as Packable>::Packed>()
     }
 }

datafusion/functions-aggregate/src/count.rs

@@ -271,6 +271,43 @@ fn get_count_accumulator(data_type: &DataType) -> Box<dyn Accumulator> {
     }
 }
 
+#[cold]
+#[inline(never)]
+fn create_distinct_groups_accumulator(
+    data_type: &DataType,
+) -> Result<Box<dyn GroupsAccumulator>> {
+    match data_type {
+        DataType::Int8 => Ok(Box::new(
+            PrimitiveDistinctCountGroupsAccumulator::<Int8Type>::new(),
+        )),
+        DataType::Int16 => Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<
+            Int16Type,
+        >::new())),
+        DataType::Int32 => Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<
+            Int32Type,
+        >::new())),
+        DataType::Int64 => Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<
+            Int64Type,
+        >::new())),
+        DataType::UInt8 => Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<
+            UInt8Type,
+        >::new())),
+        DataType::UInt16 => Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<
+            UInt16Type,
+        >::new())),
+        DataType::UInt32 => Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<
+            UInt32Type,
+        >::new())),
+        DataType::UInt64 => Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<
+            UInt64Type,
+        >::new())),
+        _ => not_impl_err!(
+            "GroupsAccumulator not supported for COUNT(DISTINCT) with {}",
+            data_type
+        ),
+    }
+}
+
 impl AggregateUDFImpl for Count {
     fn name(&self) -> &str {
         "count"
@@ -340,69 +377,30 @@ impl AggregateUDFImpl for Count {
         if args.exprs.len() != 1 {
             return false;
         }
-        if args.is_distinct {
-            // Only support primitive integer types for now
-            matches!(
-                args.expr_fields[0].data_type(),
-                DataType::Int8
-                    | DataType::Int16
-                    | DataType::Int32
-                    | DataType::Int64
-                    | DataType::UInt8
-                    | DataType::UInt16
-                    | DataType::UInt32
-                    | DataType::UInt64
-            )
-        } else {
-            true
+        if !args.is_distinct {
+            return true;
         }
+        matches!(
+            args.expr_fields[0].data_type(),
+            DataType::Int8
+                | DataType::Int16
+                | DataType::Int32
+                | DataType::Int64
+                | DataType::UInt8
+                | DataType::UInt16
+                | DataType::UInt32
+                | DataType::UInt64
+        )
     }
 
     fn create_groups_accumulator(
         &self,
         args: AccumulatorArgs,
     ) -> Result<Box<dyn GroupsAccumulator>> {
-        if args.is_distinct {
-            let data_type = args.expr_fields[0].data_type();
-            return match data_type {
-                DataType::Int8 => Ok(Box::new(
-                    PrimitiveDistinctCountGroupsAccumulator::<Int8Type>::new(),
-                )),
-                DataType::Int16 => Ok(Box::new(
-                    PrimitiveDistinctCountGroupsAccumulator::<Int16Type>::new(),
-                )),
-                DataType::Int32 => Ok(Box::new(
-                    PrimitiveDistinctCountGroupsAccumulator::<Int32Type>::new(),
-                )),
-                DataType::Int64 => Ok(Box::new(
-                    PrimitiveDistinctCountGroupsAccumulator::<Int64Type>::new(),
-                )),
-                DataType::UInt8 => Ok(Box::new(
-                    PrimitiveDistinctCountGroupsAccumulator::<UInt8Type>::new(),
-                )),
-                DataType::UInt16 => {
-                    Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<
-                        UInt16Type,
-                    >::new()))
-                }
-                DataType::UInt32 => {
-                    Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<
-                        UInt32Type,
-                    >::new()))
-                }
-                DataType::UInt64 => {
-                    Ok(Box::new(PrimitiveDistinctCountGroupsAccumulator::<
-                        UInt64Type,
-                    >::new()))
-                }
-                _ => not_impl_err!(
-                    "GroupsAccumulator not supported for COUNT(DISTINCT) with {}",
-                    data_type
-                ),
-            };
+        if !args.is_distinct {
+            return Ok(Box::new(CountGroupsAccumulator::new()));
         }
-        // instantiate specialized accumulator
-        Ok(Box::new(CountGroupsAccumulator::new()))
+        create_distinct_groups_accumulator(args.expr_fields[0].data_type())
     }
 
     fn reverse_expr(&self) -> ReversedUDAF {