diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 1746edd8239f2..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()), @@ -841,6 +835,9 @@ impl AggregateUDFImpl for ApproxDistinct { | DataType::Struct(_) | DataType::Union(_, _) | 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)))) } @@ -876,10 +873,40 @@ 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 matches!(data_type, DataType::Dictionary(_, _)) { + return is_supported_type(data_type); + } + matches!( data_type, DataType::UInt32 @@ -927,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 6ee38c48f5f1b..0e8d1eec80eb1 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -1793,6 +1793,67 @@ 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 + +# 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 CREATE TABLE approx_distinct_list_test (g INT, l INT[]) AS VALUES