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 c7b466d4f0e0c..00c1a47b9eafb 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/native.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/native.rs @@ -26,6 +26,7 @@ use std::hash::Hash; use std::mem::size_of_val; use std::sync::Arc; +use arrow::array::Array; use arrow::array::ArrayRef; use arrow::array::BooleanArray; use arrow::array::PrimitiveArray; @@ -86,11 +87,15 @@ where } let arr = as_primitive_array::(&values[0])?; - arr.iter().for_each(|value| { - if let Some(value) = value { + if arr.null_count() == 0 { + // Fast path: no nulls, so skip the per-element validity check and + // insert directly from the values buffer (mirrors `merge_batch`). + self.values.extend(arr.values().iter().copied()); + } else { + arr.iter().flatten().for_each(|value| { self.values.insert(value); - } - }); + }); + } Ok(()) } @@ -617,3 +622,42 @@ impl Accumulator for BooleanDistinctCountAccumulator { size_of_val(self) } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int64Array; + use arrow::datatypes::Int64Type; + + #[test] + fn update_batch_null_free_fast_path_agrees_with_general_path() { + // The null-free fast path must produce the same distinct set as the + // general (validity-checking) path. + let dense: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3, 2, 1])); + let sparse: ArrayRef = Arc::new(Int64Array::from(vec![ + Some(1), + None, + Some(2), + None, + Some(3), + Some(2), + Some(1), + ])); + + let mut dense_acc = + PrimitiveDistinctCountAccumulator::::new(&DataType::Int64); + dense_acc + .update_batch(std::slice::from_ref(&dense)) + .unwrap(); + + let mut sparse_acc = + PrimitiveDistinctCountAccumulator::::new(&DataType::Int64); + sparse_acc + .update_batch(std::slice::from_ref(&sparse)) + .unwrap(); + + // Both should count the 3 distinct non-null values {1, 2, 3}. + assert_eq!(dense_acc.evaluate().unwrap(), ScalarValue::Int64(Some(3))); + assert_eq!(sparse_acc.evaluate().unwrap(), ScalarValue::Int64(Some(3))); + } +}