From d21556c429d9f048ef929f9564c7b3de4ba4173d Mon Sep 17 00:00:00 2001 From: Michael Kleen Date: Tue, 18 Aug 2026 09:44:57 +0200 Subject: [PATCH 1/4] Add support for dictionary for approx_distinct --- .../src/approx_distinct.rs | 2 ++ .../sqllogictest/test_files/aggregate.slt | 34 +++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 1746edd8239f2..00a69713b79db 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -840,6 +840,7 @@ impl AggregateUDFImpl for ApproxDistinct { | DataType::Map(_, _) | DataType::Struct(_) | DataType::Union(_, _) + | DataType::Dictionary(_, _) | DataType::LargeBinary => Box::new(HLLAccumulator::new()), DataType::Null => { Box::new(NoopAccumulator::new(ScalarValue::UInt64(Some(0)))) @@ -919,6 +920,7 @@ fn is_hll_groups_type(data_type: &DataType) -> bool { | DataType::Map(_, _) | DataType::Struct(_) | DataType::Union(_, _) + | DataType::Dictionary(_, _) ) } diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 6ee38c48f5f1b..5f49a32e5b528 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -1793,6 +1793,40 @@ SELECT g, approx_distinct(arrow_cast(arrow_cast(s, 'Binary'), 'FixedSizeBinary(1 4 1 +# 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 + + # List statement ok CREATE TABLE approx_distinct_list_test (g INT, l INT[]) AS VALUES From 5f607294ffa758bf1d8bc16ae29fe4c674ee3cc7 Mon Sep 17 00:00:00 2001 From: Michael Kleen Date: Fri, 4 Sep 2026 03:43:44 +0200 Subject: [PATCH 2/4] Check if value_type is group type for dictionary --- datafusion/functions-aggregate/src/approx_distinct.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 00a69713b79db..4489e5aca4612 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -881,6 +881,11 @@ impl AggregateUDFImpl for ApproxDistinct { /// [`HllGroupsAccumulator`]. The fixed-domain types (booleans / small ints) and /// `Null` fall back to the per-group [`Accumulator`] path. fn is_hll_groups_type(data_type: &DataType) -> bool { + + if let DataType::Dictionary(_, value_type) = data_type { + return is_hll_groups_type(value_type); + } + matches!( data_type, DataType::UInt32 @@ -920,7 +925,6 @@ fn is_hll_groups_type(data_type: &DataType) -> bool { | DataType::Map(_, _) | DataType::Struct(_) | DataType::Union(_, _) - | DataType::Dictionary(_, _) ) } From 9f9aa298b1574138a2c3932097e58631e99d8e77 Mon Sep 17 00:00:00 2001 From: Michael Kleen Date: Fri, 4 Sep 2026 04:38:57 +0200 Subject: [PATCH 3/4] Fix fmt --- datafusion/functions-aggregate/src/approx_distinct.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 4489e5aca4612..04c78484f638d 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -881,7 +881,6 @@ impl AggregateUDFImpl for ApproxDistinct { /// [`HllGroupsAccumulator`]. The fixed-domain types (booleans / small ints) and /// `Null` fall back to the per-group [`Accumulator`] path. fn is_hll_groups_type(data_type: &DataType) -> bool { - if let DataType::Dictionary(_, value_type) = data_type { return is_hll_groups_type(value_type); } From e5132c71c7da6d21f607eea512776c48f0f90a8f Mon Sep 17 00:00:00 2001 From: Michael Kleen Date: Fri, 4 Sep 2026 09:48:39 +0200 Subject: [PATCH 4/4] Validate dictionary value type for approx_distinct --- .../src/approx_distinct.rs | 94 ++++++++++++++++--- .../sqllogictest/test_files/aggregate.slt | 27 ++++++ 2 files changed, 108 insertions(+), 13 deletions(-) diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 04c78484f638d..672150ee9b67e 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -735,11 +735,9 @@ impl AggregateUDFImpl for ApproxDistinct { ) .into(), ]), - DataType::Boolean - | DataType::UInt8 - | DataType::Int8 - | DataType::UInt16 - | DataType::Int16 => get_fixed_domain_state_field(args.name, data_type), + _ if is_fixed_domain_type(data_type) => { + get_fixed_domain_state_field(args.name, data_type) + } _ => Ok(vec![ Field::new( format_state_name(args.name, "hll_registers"), @@ -756,11 +754,7 @@ impl AggregateUDFImpl for ApproxDistinct { // For primitive types, use specialized accumulators for better performance. let accumulator: Box = match data_type { - DataType::Boolean - | DataType::UInt8 - | DataType::Int8 - | DataType::UInt16 - | DataType::Int16 => { + _ if is_fixed_domain_type(data_type) => { return get_fixed_domain_approx_accumulator(data_type); } DataType::UInt32 => Box::new(NumericHLLAccumulator::::new()), @@ -840,8 +834,10 @@ impl AggregateUDFImpl for ApproxDistinct { | DataType::Map(_, _) | DataType::Struct(_) | DataType::Union(_, _) - | DataType::Dictionary(_, _) | DataType::LargeBinary => Box::new(HLLAccumulator::new()), + DataType::Dictionary(_, _) if is_supported_type(data_type) => { + Box::new(HLLAccumulator::new()) + } DataType::Null => { Box::new(NoopAccumulator::new(ScalarValue::UInt64(Some(0)))) } @@ -877,12 +873,38 @@ impl AggregateUDFImpl for ApproxDistinct { } } +fn is_fixed_domain_type(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Boolean + | 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 +} + /// Returns true for the data types backed by the HyperLogLog /// [`HllGroupsAccumulator`]. The fixed-domain types (booleans / small ints) and /// `Null` fall back to the per-group [`Accumulator`] path. fn is_hll_groups_type(data_type: &DataType) -> bool { - if let DataType::Dictionary(_, value_type) = data_type { - return is_hll_groups_type(value_type); + if matches!(data_type, DataType::Dictionary(_, _)) { + return is_supported_type(data_type); } matches!( @@ -932,6 +954,52 @@ mod tests { use super::*; use std::hash::BuildHasher; + #[test] + fn dictionary_support() { + for value_type in [ + DataType::Boolean, + 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)); + } + } + #[cfg(not(feature = "force_hash_collisions"))] mod real_hash_test { use super::*; diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 5f49a32e5b528..0e8d1eec80eb1 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -1826,6 +1826,33 @@ SELECT g, approx_distinct(arrow_cast(i, 'Dictionary(Int64, Int32)')) FROM approx 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; + # List statement ok