perf : Optimize count distinct using bitmaps instead of hashsets for smaller datatypes (#21456)

## Which issue does this PR close?

Remove hashset based accumulators for smaller int data types and use
bitmaps. Follow up of : #21453
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->

- Closes #21488

## Rationale for this change

<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->

## What changes are included in this PR?

<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->

## Are these changes tested?

<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code

If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->

## Are there any user-facing changes?

<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->

<!--
If there are any breaking changes to public APIs, please add the `api
change` label.
-->

---------

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

Author: coderfender

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

@@ -22,5 +22,9 @@ mod native;
 pub use bytes::BytesDistinctCountAccumulator;
 pub use bytes::BytesViewDistinctCountAccumulator;
 pub use dict::DictionaryCountAccumulator;
+pub use native::Bitmap65536DistinctCountAccumulator;
+pub use native::Bitmap65536DistinctCountAccumulatorI16;
+pub use native::BoolArray256DistinctCountAccumulator;
+pub use native::BoolArray256DistinctCountAccumulatorI8;
 pub use native::FloatDistinctCountAccumulator;
 pub use native::PrimitiveDistinctCountAccumulator;

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)
     }
@@ -165,3 +167,354 @@ impl<T: ArrowPrimitiveType + Debug> Accumulator for FloatDistinctCountAccumulato
         size_of_val(self) + self.values.size()
     }
 }
+
+/// Optimized COUNT DISTINCT accumulator for u8 using a bool array.
+/// Uses 256 bytes to track all possible u8 values.
+#[derive(Debug)]
+pub struct BoolArray256DistinctCountAccumulator {
+    seen: [bool; 256],
+}
+
+impl BoolArray256DistinctCountAccumulator {
+    pub fn new() -> Self {
+        Self { seen: [false; 256] }
+    }
+
+    #[inline]
+    fn count(&self) -> i64 {
+        self.seen.iter().filter(|&&b| b).count() as i64
+    }
+}
+
+impl Default for BoolArray256DistinctCountAccumulator {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Accumulator for BoolArray256DistinctCountAccumulator {
+    #[inline(never)]
+    fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
+        if values.is_empty() {
+            return Ok(());
+        }
+
+        let arr = as_primitive_array::<arrow::datatypes::UInt8Type>(&values[0])?;
+        for value in arr.iter().flatten() {
+            self.seen[value as usize] = true;
+        }
+        Ok(())
+    }
+
+    fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> {
+        if states.is_empty() {
+            return Ok(());
+        }
+
+        let arr = as_list_array(&states[0])?;
+        arr.iter().try_for_each(|maybe_list| {
+            if let Some(list) = maybe_list {
+                let list = as_primitive_array::<arrow::datatypes::UInt8Type>(&list)?;
+                for value in list.values().iter() {
+                    self.seen[*value as usize] = true;
+                }
+            };
+            Ok(())
+        })
+    }
+
+    fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
+        let values: Vec<u8> = self
+            .seen
+            .iter()
+            .enumerate()
+            .filter_map(|(idx, &seen)| if seen { Some(idx as u8) } else { None })
+            .collect();
+
+        let arr = Arc::new(
+            PrimitiveArray::<arrow::datatypes::UInt8Type>::from_iter_values(values),
+        );
+        Ok(vec![
+            SingleRowListArrayBuilder::new(arr).build_list_scalar(),
+        ])
+    }
+
+    fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
+        Ok(ScalarValue::Int64(Some(self.count())))
+    }
+
+    fn size(&self) -> usize {
+        size_of_val(self) + 256
+    }
+}
+
+/// Optimized COUNT DISTINCT accumulator for i8 using a bool array.
+/// Uses 256 bytes to track all possible i8 values (mapped to 0..255).
+#[derive(Debug)]
+pub struct BoolArray256DistinctCountAccumulatorI8 {
+    seen: [bool; 256],
+}
+
+impl BoolArray256DistinctCountAccumulatorI8 {
+    pub fn new() -> Self {
+        Self { seen: [false; 256] }
+    }
+
+    #[inline]
+    fn count(&self) -> i64 {
+        self.seen.iter().filter(|&&b| b).count() as i64
+    }
+}
+
+impl Default for BoolArray256DistinctCountAccumulatorI8 {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Accumulator for BoolArray256DistinctCountAccumulatorI8 {
+    #[inline(never)]
+    fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
+        if values.is_empty() {
+            return Ok(());
+        }
+
+        let arr = as_primitive_array::<arrow::datatypes::Int8Type>(&values[0])?;
+        for value in arr.iter().flatten() {
+            self.seen[value as u8 as usize] = true;
+        }
+        Ok(())
+    }
+
+    fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> {
+        if states.is_empty() {
+            return Ok(());
+        }
+
+        let arr = as_list_array(&states[0])?;
+        arr.iter().try_for_each(|maybe_list| {
+            if let Some(list) = maybe_list {
+                let list = as_primitive_array::<arrow::datatypes::Int8Type>(&list)?;
+                for value in list.values().iter() {
+                    self.seen[*value as u8 as usize] = true;
+                }
+            };
+            Ok(())
+        })
+    }
+
+    fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
+        let values: Vec<i8> = self
+            .seen
+            .iter()
+            .enumerate()
+            .filter_map(
+                |(idx, &seen)| {
+                    if seen { Some(idx as u8 as i8) } else { None }
+                },
+            )
+            .collect();
+
+        let arr = Arc::new(
+            PrimitiveArray::<arrow::datatypes::Int8Type>::from_iter_values(values),
+        );
+        Ok(vec![
+            SingleRowListArrayBuilder::new(arr).build_list_scalar(),
+        ])
+    }
+
+    fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
+        Ok(ScalarValue::Int64(Some(self.count())))
+    }
+
+    fn size(&self) -> usize {
+        size_of_val(self) + 256
+    }
+}
+
+/// Optimized COUNT DISTINCT accumulator for u16 using a 65536-bit bitmap.
+/// Uses 8KB (1024 x u64) to track all possible u16 values.
+#[derive(Debug)]
+pub struct Bitmap65536DistinctCountAccumulator {
+    bitmap: Box<[u64; 1024]>,
+}
+
+impl Bitmap65536DistinctCountAccumulator {
+    pub fn new() -> Self {
+        Self {
+            bitmap: Box::new([0; 1024]),
+        }
+    }
+
+    #[inline]
+    fn set_bit(&mut self, value: u16) {
+        let word = (value / 64) as usize;
+        let bit = value % 64;
+        self.bitmap[word] |= 1u64 << bit;
+    }
+
+    #[inline]
+    fn count(&self) -> i64 {
+        self.bitmap.iter().map(|w| w.count_ones() as i64).sum()
+    }
+}
+
+impl Default for Bitmap65536DistinctCountAccumulator {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Accumulator for Bitmap65536DistinctCountAccumulator {
+    #[inline(never)]
+    fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
+        if values.is_empty() {
+            return Ok(());
+        }
+
+        let arr = as_primitive_array::<arrow::datatypes::UInt16Type>(&values[0])?;
+        for value in arr.iter().flatten() {
+            self.set_bit(value);
+        }
+        Ok(())
+    }
+
+    fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> {
+        if states.is_empty() {
+            return Ok(());
+        }
+
+        let arr = as_list_array(&states[0])?;
+        arr.iter().try_for_each(|maybe_list| {
+            if let Some(list) = maybe_list {
+                let list = as_primitive_array::<arrow::datatypes::UInt16Type>(&list)?;
+                for value in list.values().iter() {
+                    self.set_bit(*value);
+                }
+            };
+            Ok(())
+        })
+    }
+
+    fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
+        let mut values = Vec::new();
+        for (word_idx, &word) in self.bitmap.iter().enumerate() {
+            if word != 0 {
+                for bit in 0..64 {
+                    if (word & (1u64 << bit)) != 0 {
+                        values.push((word_idx as u16) * 64 + bit);
+                    }
+                }
+            }
+        }
+
+        let arr = Arc::new(
+            PrimitiveArray::<arrow::datatypes::UInt16Type>::from_iter_values(values),
+        );
+        Ok(vec![
+            SingleRowListArrayBuilder::new(arr).build_list_scalar(),
+        ])
+    }
+
+    fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
+        Ok(ScalarValue::Int64(Some(self.count())))
+    }
+
+    fn size(&self) -> usize {
+        size_of_val(self) + 8192
+    }
+}
+
+/// Optimized COUNT DISTINCT accumulator for i16 using a 65536-bit bitmap.
+/// Uses 8KB (1024 x u64) to track all possible i16 values (mapped to 0..65535).
+#[derive(Debug)]
+pub struct Bitmap65536DistinctCountAccumulatorI16 {
+    bitmap: Box<[u64; 1024]>,
+}
+
+impl Bitmap65536DistinctCountAccumulatorI16 {
+    pub fn new() -> Self {
+        Self {
+            bitmap: Box::new([0; 1024]),
+        }
+    }
+
+    #[inline]
+    fn set_bit(&mut self, value: i16) {
+        let idx = value as u16;
+        let word = (idx / 64) as usize;
+        let bit = idx % 64;
+        self.bitmap[word] |= 1u64 << bit;
+    }
+
+    #[inline]
+    fn count(&self) -> i64 {
+        self.bitmap.iter().map(|w| w.count_ones() as i64).sum()
+    }
+}
+
+impl Default for Bitmap65536DistinctCountAccumulatorI16 {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Accumulator for Bitmap65536DistinctCountAccumulatorI16 {
+    #[inline(never)]
+    fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
+        if values.is_empty() {
+            return Ok(());
+        }
+
+        let arr = as_primitive_array::<arrow::datatypes::Int16Type>(&values[0])?;
+        for value in arr.iter().flatten() {
+            self.set_bit(value);
+        }
+        Ok(())
+    }
+
+    fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> {
+        if states.is_empty() {
+            return Ok(());
+        }
+
+        let arr = as_list_array(&states[0])?;
+        arr.iter().try_for_each(|maybe_list| {
+            if let Some(list) = maybe_list {
+                let list = as_primitive_array::<arrow::datatypes::Int16Type>(&list)?;
+                for value in list.values().iter() {
+                    self.set_bit(*value);
+                }
+            };
+            Ok(())
+        })
+    }
+
+    fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
+        let mut values = Vec::new();
+        for (word_idx, &word) in self.bitmap.iter().enumerate() {
+            if word != 0 {
+                for bit in 0..64 {
+                    if (word & (1u64 << bit)) != 0 {
+                        values.push(((word_idx as u16) * 64 + bit) as i16);
+                    }
+                }
+            }
+        }
+
+        let arr = Arc::new(
+            PrimitiveArray::<arrow::datatypes::Int16Type>::from_iter_values(values),
+        );
+        Ok(vec![
+            SingleRowListArrayBuilder::new(arr).build_list_scalar(),
+        ])
+    }
+
+    fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
+        Ok(ScalarValue::Int64(Some(self.count())))
+    }
+
+    fn size(&self) -> usize {
+        size_of_val(self) + 8192
+    }
+}