You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by bl...@apache.org on 2025/06/09 20:11:57 UTC
(datafusion) branch main updated: Fix distinct count for DictionaryArray to correctly account for nulls in values array (#16258)
This is an automated email from the ASF dual-hosted git repository.
blaginin 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 bd85bed4b5 Fix distinct count for DictionaryArray to correctly account for nulls in values array (#16258)
bd85bed4b5 is described below
commit bd85bed4b50a169948060e9aade3e876f0cb702b
Author: kosiew <ko...@gmail.com>
AuthorDate: Tue Jun 10 04:11:52 2025 +0800
Fix distinct count for DictionaryArray to correctly account for nulls in values array (#16258)
* Add tests
* trigger ci
* Update tpch, clickbench, sort_tpch to mark failed queries (#16182)
* Move struct QueryResult to util/run.rs
* Modify benches to continue query execution even on failure
* Mark benchmark query success on output json
* Adjust slttest to pass without RUST_BACKTRACE enabled (#16251)
* add more tests where the dict keys are not null but dict values are null
* Revert "add more tests where the dict keys are not null but dict values are null"
This reverts commit c745daeef45aac172aa57b4d5353a4544a45cd2f.
* Add tests for count and count_distinct with dictionary arrays containing null values
* resolve merge conflict, reorder imports
* Add helper function to create dictionary array with non-null keys and null values
* Add test for count_distinct accumulator with dictionary array containing all null values
* add tests to aggregate.rs
* remove redundant comments in get_formatted_results function
* remove redundant safety checks in count distinct dictionary test
* refactor: introduce helper function for output assertion
* test: add count distinct dictionary handling for null values
* refactor: streamline imports and improve code organization in count.rs
* fix: add missing import for batches_to_string in aggregates.rs
* test: reorganize tests
* refactor: simplify COUNT(DISTINCT) tests and remove unused helper functions
* refactor: trim redundant tests
---------
Co-authored-by: ding-young <ls...@snu.ac.kr>
Co-authored-by: Andrew Lamb <an...@nerdnetworks.org>
Co-authored-by: Dmitrii Blaginin <dm...@blaginin.me>
---
datafusion/core/tests/sql/aggregates.rs | 121 +++++++++++++++++-
datafusion/functions-aggregate/src/count.rs | 151 ++++++++++++++++-------
datafusion/sqllogictest/test_files/aggregate.slt | 14 +++
3 files changed, 239 insertions(+), 47 deletions(-)
diff --git a/datafusion/core/tests/sql/aggregates.rs b/datafusion/core/tests/sql/aggregates.rs
index 52372e01d4..7a5868874b 100644
--- a/datafusion/core/tests/sql/aggregates.rs
+++ b/datafusion/core/tests/sql/aggregates.rs
@@ -16,7 +16,9 @@
// under the License.
use super::*;
-use datafusion::scalar::ScalarValue;
+use datafusion::common::test_util::batches_to_string;
+use datafusion_catalog::MemTable;
+use datafusion_common::ScalarValue;
#[tokio::test]
async fn csv_query_array_agg_distinct() -> Result<()> {
@@ -321,3 +323,120 @@ async fn test_accumulator_row_accumulator() -> Result<()> {
Ok(())
}
+
+/// Test that COUNT(DISTINCT) correctly handles dictionary arrays with all null values.
+/// Verifies behavior across both single and multiple partitions.
+#[tokio::test]
+async fn count_distinct_dictionary_all_null_values() -> Result<()> {
+ let n: usize = 5;
+ let num = Arc::new(Int32Array::from_iter(0..n as i32)) as ArrayRef;
+
+ // Create dictionary where all indices point to a null value (index 0)
+ let dict_values = StringArray::from(vec![None, Some("abc")]);
+ let dict_indices = Int32Array::from(vec![0; n]);
+ let dict = DictionaryArray::new(dict_indices, Arc::new(dict_values));
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("num1", DataType::Int32, false),
+ Field::new("num2", DataType::Int32, false),
+ Field::new(
+ "dict",
+ DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
+ true,
+ ),
+ ]));
+
+ let batch = RecordBatch::try_new(
+ schema.clone(),
+ vec![num.clone(), num.clone(), Arc::new(dict)],
+ )?;
+
+ // Test with single partition
+ let ctx =
+ SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1));
+ let provider = MemTable::try_new(schema.clone(), vec![vec![batch.clone()]])?;
+ ctx.register_table("t", Arc::new(provider))?;
+
+ let df = ctx
+ .sql("SELECT count(distinct dict) as cnt, count(num2) FROM t GROUP BY num1")
+ .await?;
+ let results = df.collect().await?;
+
+ assert_snapshot!(
+ batches_to_string(&results),
+ @r###"
+ +-----+---------------+
+ | cnt | count(t.num2) |
+ +-----+---------------+
+ | 0 | 1 |
+ | 0 | 1 |
+ | 0 | 1 |
+ | 0 | 1 |
+ | 0 | 1 |
+ +-----+---------------+
+ "###
+ );
+
+ // Test with multiple partitions
+ let ctx_multi =
+ SessionContext::new_with_config(SessionConfig::new().with_target_partitions(2));
+ let provider_multi = MemTable::try_new(schema, vec![vec![batch]])?;
+ ctx_multi.register_table("t", Arc::new(provider_multi))?;
+
+ let df_multi = ctx_multi
+ .sql("SELECT count(distinct dict) as cnt, count(num2) FROM t GROUP BY num1")
+ .await?;
+ let results_multi = df_multi.collect().await?;
+
+ // Results should be identical across partition configurations
+ assert_eq!(
+ batches_to_string(&results),
+ batches_to_string(&results_multi)
+ );
+
+ Ok(())
+}
+
+/// Test COUNT(DISTINCT) with mixed null and non-null dictionary values
+#[tokio::test]
+async fn count_distinct_dictionary_mixed_values() -> Result<()> {
+ let n: usize = 6;
+ let num = Arc::new(Int32Array::from_iter(0..n as i32)) as ArrayRef;
+
+ // Dictionary values array with nulls and non-nulls
+ let dict_values = StringArray::from(vec![None, Some("abc"), Some("def"), None]);
+ // Create indices that point to both null and non-null values
+ let dict_indices = Int32Array::from(vec![0, 1, 2, 0, 1, 3]);
+ let dict = DictionaryArray::new(dict_indices, Arc::new(dict_values));
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("num1", DataType::Int32, false),
+ Field::new(
+ "dict",
+ DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
+ true,
+ ),
+ ]));
+
+ let batch = RecordBatch::try_new(schema.clone(), vec![num, Arc::new(dict)])?;
+ let provider = MemTable::try_new(schema, vec![vec![batch]])?;
+ let ctx = SessionContext::new();
+ ctx.register_table("t", Arc::new(provider))?;
+
+ // COUNT(DISTINCT) should only count non-null values "abc" and "def"
+ let df = ctx.sql("SELECT count(distinct dict) FROM t").await?;
+ let results = df.collect().await?;
+
+ assert_snapshot!(
+ batches_to_string(&results),
+ @r###"
+ +------------------------+
+ | count(DISTINCT t.dict) |
+ +------------------------+
+ | 2 |
+ +------------------------+
+ "###
+ );
+
+ Ok(())
+}
diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs
index 0d10b657bc..d1fe410321 100644
--- a/datafusion/functions-aggregate/src/count.rs
+++ b/datafusion/functions-aggregate/src/count.rs
@@ -16,54 +16,49 @@
// under the License.
use ahash::RandomState;
-use datafusion_common::stats::Precision;
-use datafusion_expr::expr::WindowFunction;
-use datafusion_functions_aggregate_common::aggregate::count_distinct::BytesViewDistinctCountAccumulator;
-use datafusion_macros::user_doc;
-use datafusion_physical_expr::expressions;
-use std::collections::HashSet;
-use std::fmt::Debug;
-use std::mem::{size_of, size_of_val};
-use std::ops::BitAnd;
-use std::sync::Arc;
-
use arrow::{
- array::{ArrayRef, AsArray},
+ array::{Array, ArrayRef, AsArray, BooleanArray, Int64Array, PrimitiveArray},
+ buffer::BooleanBuffer,
compute,
datatypes::{
DataType, Date32Type, Date64Type, Decimal128Type, Decimal256Type, Field,
- Float16Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type,
- Time32MillisecondType, Time32SecondType, Time64MicrosecondType,
+ FieldRef, Float16Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type,
+ Int8Type, Time32MillisecondType, Time32SecondType, Time64MicrosecondType,
Time64NanosecondType, TimeUnit, TimestampMicrosecondType,
TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType,
UInt16Type, UInt32Type, UInt64Type, UInt8Type,
},
};
-
-use arrow::datatypes::FieldRef;
-use arrow::{
- array::{Array, BooleanArray, Int64Array, PrimitiveArray},
- buffer::BooleanBuffer,
-};
use datafusion_common::{
- downcast_value, internal_err, not_impl_err, Result, ScalarValue,
-};
-use datafusion_expr::function::StateFieldsArgs;
-use datafusion_expr::{
- function::AccumulatorArgs, utils::format_state_name, Accumulator, AggregateUDFImpl,
- Documentation, EmitTo, GroupsAccumulator, SetMonotonicity, Signature, Volatility,
+ downcast_value, internal_err, not_impl_err, stats::Precision,
+ utils::expr::COUNT_STAR_EXPANSION, Result, ScalarValue,
};
use datafusion_expr::{
- Expr, ReversedUDAF, StatisticsArgs, TypeSignature, WindowFunctionDefinition,
+ expr::WindowFunction,
+ function::{AccumulatorArgs, StateFieldsArgs},
+ utils::format_state_name,
+ Accumulator, AggregateUDFImpl, Documentation, EmitTo, Expr, GroupsAccumulator,
+ ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs, TypeSignature, Volatility,
+ WindowFunctionDefinition,
};
-use datafusion_functions_aggregate_common::aggregate::count_distinct::{
- BytesDistinctCountAccumulator, DictionaryCountAccumulator,
- FloatDistinctCountAccumulator, PrimitiveDistinctCountAccumulator,
+use datafusion_functions_aggregate_common::aggregate::{
+ count_distinct::BytesDistinctCountAccumulator,
+ count_distinct::BytesViewDistinctCountAccumulator,
+ count_distinct::DictionaryCountAccumulator,
+ count_distinct::FloatDistinctCountAccumulator,
+ count_distinct::PrimitiveDistinctCountAccumulator,
+ groups_accumulator::accumulate::accumulate_indices,
};
-use datafusion_functions_aggregate_common::aggregate::groups_accumulator::accumulate::accumulate_indices;
+use datafusion_macros::user_doc;
+use datafusion_physical_expr::expressions;
use datafusion_physical_expr_common::binary_map::OutputType;
-
-use datafusion_common::utils::expr::COUNT_STAR_EXPANSION;
+use std::{
+ collections::HashSet,
+ fmt::Debug,
+ mem::{size_of, size_of_val},
+ ops::BitAnd,
+ sync::Arc,
+};
make_udaf_expr_and_func!(
Count,
count,
@@ -712,8 +707,8 @@ impl Accumulator for DistinctCountAccumulator {
}
(0..arr.len()).try_for_each(|index| {
- if !arr.is_null(index) {
- let scalar = ScalarValue::try_from_array(arr, index)?;
+ let scalar = ScalarValue::try_from_array(arr, index)?;
+ if !scalar.is_null() {
self.values.insert(scalar);
}
Ok(())
@@ -758,14 +753,28 @@ impl Accumulator for DistinctCountAccumulator {
#[cfg(test)]
mod tests {
- use std::sync::Arc;
use super::*;
-
- use arrow::array::{Int32Array, NullArray};
- use arrow::datatypes::{DataType, Field, Int32Type, Schema};
+ use arrow::{
+ array::{DictionaryArray, Int32Array, NullArray, StringArray},
+ datatypes::{DataType, Field, Int32Type, Schema},
+ };
use datafusion_expr::function::AccumulatorArgs;
use datafusion_physical_expr::expressions::Column;
+ use std::sync::Arc;
+ /// Helper function to create a dictionary array with non-null keys but some null values
+ /// Returns a dictionary array where:
+ /// - keys are [0, 1, 2, 0, 1] (all non-null)
+ /// - values are ["a", null, "c"]
+ /// - so the keys reference: "a", null, "c", "a", null
+ fn create_dictionary_with_null_values() -> Result<DictionaryArray<Int32Type>> {
+ let values = StringArray::from(vec![Some("a"), None, Some("c")]);
+ let keys = Int32Array::from(vec![0, 1, 2, 0, 1]); // references "a", null, "c", "a", null
+ Ok(DictionaryArray::<Int32Type>::try_new(
+ keys,
+ Arc::new(values),
+ )?)
+ }
#[test]
fn count_accumulator_nulls() -> Result<()> {
@@ -803,15 +812,12 @@ mod tests {
order_bys: &[],
};
- let inner_dict = arrow::array::DictionaryArray::<Int32Type>::from_iter([
- "a", "b", "c", "d", "a", "b",
- ]);
+ let inner_dict =
+ DictionaryArray::<Int32Type>::from_iter(["a", "b", "c", "d", "a", "b"]);
let keys = Int32Array::from(vec![0, 1, 2, 0, 3, 1]);
- let dict_of_dict = arrow::array::DictionaryArray::<Int32Type>::try_new(
- keys,
- Arc::new(inner_dict),
- )?;
+ let dict_of_dict =
+ DictionaryArray::<Int32Type>::try_new(keys, Arc::new(inner_dict))?;
let mut acc = count.accumulator(args)?;
acc.update_batch(&[Arc::new(dict_of_dict)])?;
@@ -819,4 +825,57 @@ mod tests {
Ok(())
}
+
+ #[test]
+ fn count_distinct_accumulator_dictionary_with_null_values() -> Result<()> {
+ let dict_array = create_dictionary_with_null_values()?;
+
+ // The expected behavior is that count_distinct should count only non-null values
+ // which in this case are "a" and "c" (appearing as 0 and 2 in keys)
+ let mut accumulator = DistinctCountAccumulator {
+ values: HashSet::default(),
+ state_data_type: dict_array.data_type().clone(),
+ };
+
+ accumulator.update_batch(&[Arc::new(dict_array)])?;
+
+ // Should have 2 distinct non-null values ("a" and "c")
+ assert_eq!(accumulator.evaluate()?, ScalarValue::Int64(Some(2)));
+ Ok(())
+ }
+
+ #[test]
+ fn count_accumulator_dictionary_with_null_values() -> Result<()> {
+ let dict_array = create_dictionary_with_null_values()?;
+
+ // The expected behavior is that count should only count non-null values
+ let mut accumulator = CountAccumulator::new();
+
+ accumulator.update_batch(&[Arc::new(dict_array)])?;
+
+ // 5 elements in the array, of which 2 reference null values (the two 1s in the keys)
+ // So we should count 3 non-null values
+ assert_eq!(accumulator.evaluate()?, ScalarValue::Int64(Some(3)));
+ Ok(())
+ }
+
+ #[test]
+ fn count_distinct_accumulator_dictionary_all_null_values() -> Result<()> {
+ // Create a dictionary array that only contains null values
+ let dict_values = StringArray::from(vec![None, Some("abc")]);
+ let dict_indices = Int32Array::from(vec![0; 5]);
+ let dict_array =
+ DictionaryArray::<Int32Type>::try_new(dict_indices, Arc::new(dict_values))?;
+
+ let mut accumulator = DistinctCountAccumulator {
+ values: HashSet::default(),
+ state_data_type: dict_array.data_type().clone(),
+ };
+
+ accumulator.update_batch(&[Arc::new(dict_array)])?;
+
+ // All referenced values are null so count(distinct) should be 0
+ assert_eq!(accumulator.evaluate()?, ScalarValue::Int64(Some(0)));
+ Ok(())
+ }
}
diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt
index 0966d9ff05..bd9d8b1f43 100644
--- a/datafusion/sqllogictest/test_files/aggregate.slt
+++ b/datafusion/sqllogictest/test_files/aggregate.slt
@@ -5029,6 +5029,20 @@ select count(distinct column1), count(distinct column2) from dict_test group by
statement ok
drop table dict_test;
+## count distinct dictionary with null values
+statement ok
+create table dict_null_test as
+ select arrow_cast(NULL, 'Dictionary(Int32, Utf8)') as d
+ from (values (1), (2), (3), (4), (5));
+
+query I
+select count(distinct d) from dict_null_test;
+----
+0
+
+statement ok
+drop table dict_null_test;
+
# avg_duration
statement ok
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org