From 952a57399f73e2e7db0440b3dd76a324e1064358 Mon Sep 17 00:00:00 2001 From: Michael Kleen Date: Fri, 4 Sep 2026 13:38:32 +0000 Subject: [PATCH 1/3] [cherry-pick] Add support for dictionary for approx_distinct (#24646) Resolves merge conflict between HEAD's type-specific HLL accumulators and the upstream dictionary support commit. Implements dictionary handling via a DictionaryAccumulator wrapper that casts to the value type before delegating to the appropriate inner accumulator. Co-Authored-By: Claude Sonnet 4.6 --- .../src/approx_distinct.rs | 197 ++++++++++++++++++ .../sqllogictest/test_files/aggregate.slt | 72 +++++++ 2 files changed, 269 insertions(+) diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index cc42b6c22bdbe..954e7c678cac0 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -395,6 +395,19 @@ impl AggregateUDFImpl for ApproxDistinct { DataType::UInt8 | DataType::Int8 | DataType::UInt16 | DataType::Int16 => { get_small_int_state_field(args.name, data_type) } + DataType::Dictionary(_, _) if is_supported_type(data_type) => { + let value_type = dictionary_value_type(data_type); + if is_fixed_domain_type(value_type) { + get_small_int_state_field(args.name, value_type) + } else { + Ok(vec![Field::new( + format_state_name(args.name, "hll_registers"), + DataType::Binary, + false, + ) + .into()]) + } + } _ => Ok(vec![ Field::new( format_state_name(args.name, "hll_registers"), @@ -448,6 +461,11 @@ impl AggregateUDFImpl for ApproxDistinct { DataType::Utf8View => Box::new(StringViewHLLAccumulator::new()), DataType::Binary => Box::new(BinaryHLLAccumulator::::new()), DataType::LargeBinary => Box::new(BinaryHLLAccumulator::::new()), + DataType::Dictionary(_, _) if is_supported_type(data_type) => { + let value_type = dictionary_value_type(data_type).clone(); + let inner = make_approx_distinct_accumulator(&value_type)?; + Box::new(DictionaryAccumulator { inner, value_type }) + } DataType::Null => { Box::new(NoopAccumulator::new(ScalarValue::UInt64(Some(0)))) } @@ -464,3 +482,182 @@ impl AggregateUDFImpl for ApproxDistinct { self.doc() } } + +#[derive(Debug)] +struct DictionaryAccumulator { + inner: Box, + value_type: DataType, +} + +impl Accumulator for DictionaryAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + let plain = arrow::compute::cast(&values[0], &self.value_type)?; + self.inner.update_batch(&[plain]) + } + + fn evaluate(&mut self) -> Result { + self.inner.evaluate() + } + + fn size(&self) -> usize { + self.inner.size() + } + + fn state(&mut self) -> Result> { + self.inner.state() + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + self.inner.merge_batch(states) + } +} + +fn make_approx_distinct_accumulator(data_type: &DataType) -> Result> { + match data_type { + DataType::UInt8 | DataType::Int8 | DataType::UInt16 | DataType::Int16 => { + get_small_int_approx_accumulator(data_type) + } + DataType::UInt32 => Ok(Box::new(NumericHLLAccumulator::::new())), + DataType::UInt64 => Ok(Box::new(NumericHLLAccumulator::::new())), + DataType::Int32 => Ok(Box::new(NumericHLLAccumulator::::new())), + DataType::Int64 => Ok(Box::new(NumericHLLAccumulator::::new())), + DataType::Date32 => Ok(Box::new(NumericHLLAccumulator::::new())), + DataType::Date64 => Ok(Box::new(NumericHLLAccumulator::::new())), + DataType::Time32(TimeUnit::Second) => { + Ok(Box::new(NumericHLLAccumulator::::new())) + } + DataType::Time32(TimeUnit::Millisecond) => { + Ok(Box::new(NumericHLLAccumulator::::new())) + } + DataType::Time64(TimeUnit::Microsecond) => { + Ok(Box::new(NumericHLLAccumulator::::new())) + } + DataType::Time64(TimeUnit::Nanosecond) => { + Ok(Box::new(NumericHLLAccumulator::::new())) + } + DataType::Timestamp(TimeUnit::Second, _) => { + Ok(Box::new(NumericHLLAccumulator::::new())) + } + DataType::Timestamp(TimeUnit::Millisecond, _) => { + Ok(Box::new(NumericHLLAccumulator::::new())) + } + DataType::Timestamp(TimeUnit::Microsecond, _) => { + Ok(Box::new(NumericHLLAccumulator::::new())) + } + DataType::Timestamp(TimeUnit::Nanosecond, _) => { + Ok(Box::new(NumericHLLAccumulator::::new())) + } + DataType::Utf8 => Ok(Box::new(StringHLLAccumulator::::new())), + DataType::LargeUtf8 => Ok(Box::new(StringHLLAccumulator::::new())), + DataType::Utf8View => Ok(Box::new(StringViewHLLAccumulator::new())), + DataType::Binary => Ok(Box::new(BinaryHLLAccumulator::::new())), + DataType::LargeBinary => Ok(Box::new(BinaryHLLAccumulator::::new())), + DataType::Null => { + Ok(Box::new(NoopAccumulator::new(ScalarValue::UInt64(Some(0))))) + } + other => not_impl_err!( + "Support for 'approx_distinct' for data type {other} is not implemented" + ), + } +} + +fn is_fixed_domain_type(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::UInt8 | DataType::Int8 | DataType::UInt16 | DataType::Int16 + ) +} + +fn is_supported_type(data_type: &DataType) -> bool { + let value_type = dictionary_value_type(data_type); + matches!(value_type, DataType::Null) + || is_fixed_domain_type(value_type) + || is_hll_groups_type(value_type) +} + +fn dictionary_value_type(data_type: &DataType) -> &DataType { + let mut value_type = data_type; + while let DataType::Dictionary(_, inner) = value_type { + value_type = inner; + } + value_type +} + +fn is_hll_groups_type(data_type: &DataType) -> bool { + if matches!(data_type, DataType::Dictionary(_, _)) { + return is_supported_type(data_type); + } + + matches!( + data_type, + DataType::UInt32 + | DataType::UInt64 + | DataType::Int32 + | DataType::Int64 + | DataType::Date32 + | DataType::Date64 + | DataType::Time32(TimeUnit::Second) + | DataType::Time32(TimeUnit::Millisecond) + | DataType::Time64(TimeUnit::Microsecond) + | DataType::Time64(TimeUnit::Nanosecond) + | DataType::Timestamp(TimeUnit::Second, _) + | DataType::Timestamp(TimeUnit::Millisecond, _) + | DataType::Timestamp(TimeUnit::Microsecond, _) + | DataType::Timestamp(TimeUnit::Nanosecond, _) + | DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Utf8View + | DataType::Binary + | DataType::LargeBinary + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn dictionary_support() { + for value_type in [ + DataType::UInt8, + DataType::Int8, + DataType::UInt16, + DataType::Int16, + DataType::Int64, + DataType::Null, + DataType::Utf8, + DataType::Binary, + ] { + let dict_type = DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(value_type.clone()), + ); + assert!(is_hll_groups_type(&dict_type)); + } + + // Nested dictionaries resolve to the innermost value + assert!(is_hll_groups_type(&DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Utf8) + )) + ))); + + // Unsupported value types are rejected + for value_type in [DataType::Float16, DataType::Float32, DataType::Float64] { + let dict_type = DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(value_type.clone()), + ); + let nested_dict_type = DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(dict_type.clone()), + ); + assert!(!is_hll_groups_type(&value_type)); + assert!(!is_supported_type(&dict_type)); + assert!(!is_hll_groups_type(&dict_type)); + assert!(!is_hll_groups_type(&nested_dict_type)); + } + } +} diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 3e6c16e12595f..1dbcd6fa5803d 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -1836,6 +1836,78 @@ SELECT approx_distinct(c14) AS a, approx_distinct(c15) AS b, approx_distinct(arr ---- 18 60 60 60 60 + +statement ok +CREATE TABLE approx_distinct_group_test (g INT, s VARCHAR, i INT) AS VALUES + (1, 'a', 10), (1, 'a', 10), (1, 'b', 20), + (2, 'c', 30), (2, 'd', 30), (2, 'c', 40), + (3, NULL, NULL), (3, NULL, NULL), + (4, 'e', 50); + +# Dictionary: dictionary-encoded values must hash identically to the plain +# (non-dictionary) values, so the counts below match the Utf8 case above. + +# Dictionary non-grouped +query I +SELECT approx_distinct(arrow_cast(s, 'Dictionary(Int32, Utf8)')) FROM approx_distinct_group_test WHERE g = 2; +---- +2 + +# Dictionary grouped +query II +SELECT g, approx_distinct(arrow_cast(s, 'Dictionary(Int32, Utf8)')) FROM approx_distinct_group_test GROUP BY g ORDER BY g; +---- +1 2 +2 2 +3 0 +4 1 + +# Dictionary with a non-string value type (Int32), also exercising a +# larger (Int64) key type +query I +SELECT approx_distinct(arrow_cast(i, 'Dictionary(Int64, Int32)')) FROM approx_distinct_group_test WHERE g = 2; +---- +2 + +query II +SELECT g, approx_distinct(arrow_cast(i, 'Dictionary(Int64, Int32)')) FROM approx_distinct_group_test GROUP BY g ORDER BY g; +---- +1 2 +2 2 +3 0 +4 1 + +# Dictionary over a fixed-domain value type (Int8). The bitmap accumulator only +# understands its native value array, so these go through the HyperLogLog groups +# accumulator and must still match the Int32 counts above. +query I +SELECT approx_distinct(arrow_cast(arrow_cast(i, 'Int8'), 'Dictionary(Int32, Int8)')) FROM approx_distinct_group_test WHERE g = 2; +---- +2 + +query II +SELECT g, approx_distinct(arrow_cast(arrow_cast(i, 'Int8'), 'Dictionary(Int32, Int8)')) FROM approx_distinct_group_test GROUP BY g ORDER BY g; +---- +1 2 +2 2 +3 0 +4 1 + +# A dictionary is supported exactly when its value type is: floats are rejected +# just like a bare Float64 is, rather than silently reaching the HLL accumulator. +statement error DataFusion error: This feature is not implemented: Support for 'approx_distinct' for data type Float64 is not implemented +SELECT approx_distinct(arrow_cast(i, 'Float64')) FROM approx_distinct_group_test; + +statement error DataFusion error: This feature is not implemented: Support for 'approx_distinct' for data type Dictionary\(Int32, Float64\) is not implemented +SELECT approx_distinct(arrow_cast(i, 'Dictionary(Int32, Float64)')) FROM approx_distinct_group_test; + +statement error DataFusion error: This feature is not implemented: Support for 'approx_distinct' for data type Dictionary\(Int32, Float64\) is not implemented +SELECT g, approx_distinct(arrow_cast(i, 'Dictionary(Int32, Float64)')) FROM approx_distinct_group_test GROUP BY g; + + + +statement ok +DROP TABLE approx_distinct_group_test; ## This test executes the APPROX_PERCENTILE_CONT aggregation against the test ## data, asserting the estimated quantiles are ±5% their actual values. ## From d8c10ab6c4774bd28ffa814a12fd59c8712ec9de Mon Sep 17 00:00:00 2001 From: RIchard Baah Date: Fri, 4 Sep 2026 10:26:06 -0400 Subject: [PATCH 2/3] fix lint --- .../src/approx_distinct.rs | 54 ++++++++++--------- 1 file changed, 29 insertions(+), 25 deletions(-) diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 954e7c678cac0..7c50e2730579a 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -400,12 +400,14 @@ impl AggregateUDFImpl for ApproxDistinct { if is_fixed_domain_type(value_type) { get_small_int_state_field(args.name, value_type) } else { - Ok(vec![Field::new( - format_state_name(args.name, "hll_registers"), - DataType::Binary, - false, - ) - .into()]) + Ok(vec![ + Field::new( + format_state_name(args.name, "hll_registers"), + DataType::Binary, + false, + ) + .into(), + ]) } } _ => Ok(vec![ @@ -512,7 +514,9 @@ impl Accumulator for DictionaryAccumulator { } } -fn make_approx_distinct_accumulator(data_type: &DataType) -> Result> { +fn make_approx_distinct_accumulator( + data_type: &DataType, +) -> Result> { match data_type { DataType::UInt8 | DataType::Int8 | DataType::UInt16 | DataType::Int16 => { get_small_int_approx_accumulator(data_type) @@ -526,27 +530,27 @@ fn make_approx_distinct_accumulator(data_type: &DataType) -> Result { Ok(Box::new(NumericHLLAccumulator::::new())) } - DataType::Time32(TimeUnit::Millisecond) => { - Ok(Box::new(NumericHLLAccumulator::::new())) - } - DataType::Time64(TimeUnit::Microsecond) => { - Ok(Box::new(NumericHLLAccumulator::::new())) - } - DataType::Time64(TimeUnit::Nanosecond) => { - Ok(Box::new(NumericHLLAccumulator::::new())) - } + DataType::Time32(TimeUnit::Millisecond) => Ok(Box::new(NumericHLLAccumulator::< + Time32MillisecondType, + >::new())), + DataType::Time64(TimeUnit::Microsecond) => Ok(Box::new(NumericHLLAccumulator::< + Time64MicrosecondType, + >::new())), + DataType::Time64(TimeUnit::Nanosecond) => Ok(Box::new(NumericHLLAccumulator::< + Time64NanosecondType, + >::new())), DataType::Timestamp(TimeUnit::Second, _) => { Ok(Box::new(NumericHLLAccumulator::::new())) } - DataType::Timestamp(TimeUnit::Millisecond, _) => { - Ok(Box::new(NumericHLLAccumulator::::new())) - } - DataType::Timestamp(TimeUnit::Microsecond, _) => { - Ok(Box::new(NumericHLLAccumulator::::new())) - } - DataType::Timestamp(TimeUnit::Nanosecond, _) => { - Ok(Box::new(NumericHLLAccumulator::::new())) - } + DataType::Timestamp(TimeUnit::Millisecond, _) => Ok(Box::new( + NumericHLLAccumulator::::new(), + )), + DataType::Timestamp(TimeUnit::Microsecond, _) => Ok(Box::new( + NumericHLLAccumulator::::new(), + )), + DataType::Timestamp(TimeUnit::Nanosecond, _) => Ok(Box::new( + NumericHLLAccumulator::::new(), + )), DataType::Utf8 => Ok(Box::new(StringHLLAccumulator::::new())), DataType::LargeUtf8 => Ok(Box::new(StringHLLAccumulator::::new())), DataType::Utf8View => Ok(Box::new(StringViewHLLAccumulator::new())), From c0fd2c88b04dcc6cbb83eec256e8be04acd8c19e Mon Sep 17 00:00:00 2001 From: RIchard Baah Date: Fri, 4 Sep 2026 11:06:32 -0400 Subject: [PATCH 3/3] fix: add blank line after DROP TABLE in aggregate.slt to terminate SLT statement Co-Authored-By: Claude Sonnet 4.6 --- datafusion/sqllogictest/test_files/aggregate.slt | 1 + 1 file changed, 1 insertion(+) diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 1dbcd6fa5803d..fa62a32b9c126 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -1908,6 +1908,7 @@ SELECT g, approx_distinct(arrow_cast(i, 'Dictionary(Int32, Float64)')) FROM appr statement ok DROP TABLE approx_distinct_group_test; + ## This test executes the APPROX_PERCENTILE_CONT aggregation against the test ## data, asserting the estimated quantiles are ±5% their actual values. ##