You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/04/16 18:59:23 UTC
(datafusion) branch main updated: perf : Optimize count distinct using bitmaps instead of hashsets for smaller datatypes (#21456)
This is an automated email from the ASF dual-hosted git repository.
github-bot pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 1f0faf92cc perf : Optimize count distinct using bitmaps instead of hashsets for smaller datatypes (#21456)
1f0faf92cc is described below
commit 1f0faf92cc4437ddb806a027a5ea15ebd27effba
Author: Bhargava Vadlamani <11...@users.noreply.github.com>
AuthorDate: Thu Apr 16 11:56:48 2026 -0700
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 : https://github.com/apache/datafusion/pull/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 https://github.com/apache/datafusion/issues/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 <an...@nerdnetworks.org>
---
.../src/aggregate/count_distinct.rs | 4 +
.../src/aggregate/count_distinct/native.rs | 353 +++++++++++++++++++++
datafusion/functions-aggregate/src/count.rs | 42 ++-
3 files changed, 382 insertions(+), 17 deletions(-)
diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct.rs
index 25b4038229..4f26d86839 100644
--- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct.rs
+++ b/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;
diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/native.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/native.rs
index e506b4acb1..fb9cfb379a 100644
--- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/native.rs
+++ b/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
+ }
+}
diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs
index 0edfc152ff..75a7018f32 100644
--- a/datafusion/functions-aggregate/src/count.rs
+++ b/datafusion/functions-aggregate/src/count.rs
@@ -21,16 +21,16 @@ use arrow::{
compute,
datatypes::{
DataType, Date32Type, Date64Type, Decimal128Type, Decimal256Type, Field,
- FieldRef, Float16Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type,
- Int64Type, Time32MillisecondType, Time32SecondType, Time64MicrosecondType,
+ FieldRef, Float16Type, Float32Type, Float64Type, Int32Type, Int64Type,
+ Time32MillisecondType, Time32SecondType, Time64MicrosecondType,
Time64NanosecondType, TimeUnit, TimestampMicrosecondType,
TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType,
- UInt8Type, UInt16Type, UInt32Type, UInt64Type,
+ UInt32Type, UInt64Type,
},
};
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::{
@@ -42,6 +42,10 @@ use datafusion_expr::{
utils::format_state_name,
};
use datafusion_functions_aggregate_common::aggregate::{
+ count_distinct::Bitmap65536DistinctCountAccumulator,
+ count_distinct::Bitmap65536DistinctCountAccumulatorI16,
+ count_distinct::BoolArray256DistinctCountAccumulator,
+ count_distinct::BoolArray256DistinctCountAccumulatorI8,
count_distinct::BytesDistinctCountAccumulator,
count_distinct::BytesViewDistinctCountAccumulator,
count_distinct::DictionaryCountAccumulator,
@@ -170,31 +174,23 @@ impl Count {
}
fn get_count_accumulator(data_type: &DataType) -> Box<dyn Accumulator> {
match data_type {
- // try and use a specialized accumulator if possible, otherwise fall back to generic accumulator
- DataType::Int8 => Box::new(PrimitiveDistinctCountAccumulator::<Int8Type>::new(
- data_type,
- )),
- DataType::Int16 => Box::new(PrimitiveDistinctCountAccumulator::<Int16Type>::new(
- data_type,
- )),
+ // HashSet-based accumulator for larger integer types
DataType::Int32 => Box::new(PrimitiveDistinctCountAccumulator::<Int32Type>::new(
data_type,
)),
DataType::Int64 => Box::new(PrimitiveDistinctCountAccumulator::<Int64Type>::new(
data_type,
)),
- DataType::UInt8 => Box::new(PrimitiveDistinctCountAccumulator::<UInt8Type>::new(
- data_type,
- )),
- DataType::UInt16 => Box::new(
- PrimitiveDistinctCountAccumulator::<UInt16Type>::new(data_type),
- ),
DataType::UInt32 => Box::new(
PrimitiveDistinctCountAccumulator::<UInt32Type>::new(data_type),
),
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)),
@@ -270,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"
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org