move_bitmap_accumulator_cold_path

Author: B Vadlamani

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

@@ -78,6 +78,7 @@ where
         ])
     }
 
+    #[inline(never)]
     fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
         if values.is_empty() {
             return Ok(());
@@ -149,6 +150,7 @@ impl<T: ArrowPrimitiveType + Debug> Accumulator for FloatDistinctCountAccumulato
         self.values.state()
     }
 
+    #[inline(never)]
     fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
         self.values.update_batch(values)
     }

datafusion/functions-aggregate/src/count.rs

@@ -30,7 +30,7 @@ use arrow::{
 };
 use datafusion_common::hash_utils::RandomState;
 use datafusion_common::{
-    HashMap, Result, ScalarValue, downcast_value, internal_err, not_impl_err,
+    HashMap, Result, ScalarValue, downcast_value, exec_err, internal_err, not_impl_err,
     stats::Precision, utils::expr::COUNT_STAR_EXPANSION,
 };
 use datafusion_expr::{
@@ -174,11 +174,6 @@ impl Count {
 }
 fn get_count_accumulator(data_type: &DataType) -> Box<dyn Accumulator> {
     match data_type {
-        // Optimized bitmap/bool array accumulators for small integer types
-        DataType::UInt8 => Box::new(BoolArray256DistinctCountAccumulator::new()),
-        DataType::Int8 => Box::new(BoolArray256DistinctCountAccumulatorI8::new()),
-        DataType::UInt16 => Box::new(Bitmap65536DistinctCountAccumulator::new()),
-        DataType::Int16 => Box::new(Bitmap65536DistinctCountAccumulatorI16::new()),
         // HashSet-based accumulator for larger integer types
         DataType::Int32 => Box::new(PrimitiveDistinctCountAccumulator::<Int32Type>::new(
             data_type,
@@ -192,6 +187,10 @@ fn get_count_accumulator(data_type: &DataType) -> Box<dyn Accumulator> {
         DataType::UInt64 => Box::new(
             PrimitiveDistinctCountAccumulator::<UInt64Type>::new(data_type),
         ),
+        // Small int types - cold path
+        DataType::UInt8 | DataType::Int8 | DataType::UInt16 | DataType::Int16 => {
+            get_small_int_accumulator(data_type).unwrap()
+        }
         DataType::Decimal128(_, _) => Box::new(PrimitiveDistinctCountAccumulator::<
             Decimal128Type,
         >::new(data_type)),
@@ -267,6 +266,18 @@ fn get_count_accumulator(data_type: &DataType) -> Box<dyn Accumulator> {
     }
 }
 
+/// Uses optimized bitmap accumulators but separated to keep hot path small
+#[cold]
+fn get_small_int_accumulator(data_type: &DataType) -> Result<Box<dyn Accumulator>> {
+    match data_type {
+        DataType::UInt8 => Ok(Box::new(BoolArray256DistinctCountAccumulator::new())),
+        DataType::Int8 => Ok(Box::new(BoolArray256DistinctCountAccumulatorI8::new())),
+        DataType::UInt16 => Ok(Box::new(Bitmap65536DistinctCountAccumulator::new())),
+        DataType::Int16 => Ok(Box::new(Bitmap65536DistinctCountAccumulatorI16::new())),
+        _ => exec_err!("unsupported accumulator for datatype: {}", data_type),
+    }
+}
+
 impl AggregateUDFImpl for Count {
     fn name(&self) -> &str {
         "count"