From ed101e9d5c13e823bfd0573150de758f5e8ba30f Mon Sep 17 00:00:00 2001 From: Michael Kleen Date: Sun, 19 Jul 2026 11:07:59 +0200 Subject: [PATCH 1/3] feat: Support Union type in approx_distinct --- .../src/approx_distinct.rs | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 6d3ea3a8ddded..ef87d68eede35 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -844,6 +844,7 @@ impl AggregateUDFImpl for ApproxDistinct { | DataType::LargeListView(_) | DataType::Map(_, _) | DataType::Struct(_) + | DataType::Union(_, _) | DataType::LargeBinary => Box::new(HLLAccumulator::new()), DataType::Null => { Box::new(NoopAccumulator::new(ScalarValue::UInt64(Some(0)))) @@ -922,6 +923,7 @@ fn is_hll_groups_type(data_type: &DataType) -> bool { | DataType::LargeListView(_) | DataType::Map(_, _) | DataType::Struct(_) + | DataType::Union(_, _) ) } @@ -1078,6 +1080,66 @@ mod tests { assert_count_numerical_acc_and_group_acc::(decimal_256, 6); } + #[test] + fn union_support_hll_acc_and_group_acc() { + use arrow::array::{Int32Array, StringArray, UnionArray}; + use arrow::datatypes::{UnionFields, UnionMode}; + + assert!(is_hll_groups_type(&DataType::Union( + UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ], + ) + .unwrap(), + UnionMode::Dense, + ))); + + let int_array = Int32Array::from(vec![5, 10, 5]); + let str_array = StringArray::from(vec!["foo", "foo"]); + let type_ids = vec![0_i8, 1, 0, 0, 1].into(); + let offsets = vec![0, 0, 1, 2, 1].into(); + let children = vec![ + Arc::new(int_array) as ArrayRef, + Arc::new(str_array) as ArrayRef, + ]; + let union_fields = [ + (0, Arc::new(Field::new("a", DataType::Int32, true))), + (1, Arc::new(Field::new("b", DataType::Utf8, true))), + ] + .into_iter() + .collect(); + let array: ArrayRef = Arc::new( + UnionArray::try_new(union_fields, type_ids, Some(offsets), children) + .unwrap(), + ); + + let mut acc = HLLAccumulator::new(); + acc.update_batch(&[Arc::clone(&array)]).unwrap(); + let non_group_count = distinct_count(&mut acc); + + let group_indices = vec![0usize; array.len()]; + let mut group_acc = HllGroupsAccumulator::new(); + group_acc + .update_batch(std::slice::from_ref(&array), &group_indices, None, 1) + .unwrap(); + let group_count = group_acc + .evaluate(EmitTo::All) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + + assert_eq!( + non_group_count, group_count, + "non-grouped and grouped paths disagree for union type" + ); + assert_eq!(non_group_count, 3, "wrong distinct count for union type"); + } + #[test] fn interval_support_numerical_acc_and_group_acc() { let year_month: ArrayRef = From 7fa00c72501af946803bb9790291f4f9c0b54cb6 Mon Sep 17 00:00:00 2001 From: Michael Kleen Date: Wed, 22 Jul 2026 09:51:21 +0200 Subject: [PATCH 2/3] Add slt --- datafusion/sqllogictest/src/test_context.rs | 41 +++++++++++++++++++ .../sqllogictest/test_files/aggregate.slt | 30 ++++++++++++++ 2 files changed, 71 insertions(+) diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index e0aaa91ef6369..fdb04edc05101 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -187,6 +187,10 @@ impl TestContext { info!("Registering table with union column"); register_union_table(test_ctx.session_ctx()) } + "aggregate.slt" => { + info!("Registering table with union column for approx_distinct"); + register_approx_distinct_union_table(test_ctx.session_ctx()) + } "dictionary_struct.slt" => { info!("Registering table with dictionary-encoded struct column"); register_dictionary_struct_table(test_ctx.session_ctx()); @@ -593,6 +597,43 @@ fn register_union_table(ctx: &SessionContext) { ctx.register_batch("union_table", batch).unwrap(); } +fn register_approx_distinct_union_table(ctx: &SessionContext) { + let union = UnionArray::try_new( + UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("i", DataType::Int32, true), + Field::new("s", DataType::Utf8, true), + ], + ) + .unwrap(), + ScalarBuffer::from(vec![0_i8, 0, 1, 1, 0, 0, 1, 0]), + Some(ScalarBuffer::from(vec![0, 1, 0, 1, 2, 3, 2, 4])), + vec![ + Arc::new(Int32Array::from(vec![ + Some(1), + Some(1), + None, + None, + Some(5), + ])), + Arc::new(StringArray::from(vec![Some("x"), Some("y"), None])), + ], + ) + .unwrap(); + + let schema = Schema::new(vec![ + Field::new("g", DataType::Int32, false), + Field::new("u", union.data_type().clone(), false), + ]); + + let g = Arc::new(Int32Array::from(vec![1, 1, 1, 2, 2, 3, 3, 4])); + let batch = RecordBatch::try_new(Arc::new(schema), vec![g, Arc::new(union)]).unwrap(); + + ctx.register_batch("approx_distinct_union_test", batch) + .unwrap(); +} + fn register_dictionary_struct_table(ctx: &SessionContext) { // Build deduplicated struct values: 3 unique structs let names = Arc::new(StringArray::from(vec!["Alice", "Bob", "Carol"])) as ArrayRef; diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index cbb9c5d0317dc..365505e2da9e2 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -2171,6 +2171,36 @@ SELECT approx_distinct(s) FROM approx_distinct_struct_test; statement ok DROP TABLE approx_distinct_struct_test; +# Union +# `approx_distinct_union_test` (g INT, u UNION) is registered +# in test_context.rs because a union value cannot be constructed from SQL. + +# Union non-grouped +query I +SELECT approx_distinct(u) FROM approx_distinct_union_test WHERE g = 1; +---- +2 + +# Union grouped +# Group 1 -> {i:1, i:1, s:"x"}=2, +# Group 2 -> {s:"y"}=1 (NULL excluded), +# Group 3 -> all null=0, +# Group 4 -> {i:5}=1 +query II +SELECT g, approx_distinct(u) FROM approx_distinct_union_test GROUP BY g ORDER BY g; +---- +1 2 +2 1 +3 0 +4 1 + +# The non-group path must agree with the grouped path on +# the same data, and distinct union values across groups are still counted overall. +query I +SELECT approx_distinct(u) FROM approx_distinct_union_test; +---- +4 + # Integers (Int32): group 1 -> {10,20}=2, group 2 -> {30,40}=2, group 3 -> 0, group 4 -> {50}=1 query II SELECT g, approx_distinct(i) FROM approx_distinct_group_test GROUP BY g ORDER BY g; From addd02c1083f597339e5ca85efbc56ca96dacdfd Mon Sep 17 00:00:00 2001 From: Michael Kleen Date: Wed, 22 Jul 2026 09:52:25 +0200 Subject: [PATCH 3/3] remove unit test --- .../src/approx_distinct.rs | 60 ------------------- 1 file changed, 60 deletions(-) diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index ef87d68eede35..f36c658e7f385 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -1080,66 +1080,6 @@ mod tests { assert_count_numerical_acc_and_group_acc::(decimal_256, 6); } - #[test] - fn union_support_hll_acc_and_group_acc() { - use arrow::array::{Int32Array, StringArray, UnionArray}; - use arrow::datatypes::{UnionFields, UnionMode}; - - assert!(is_hll_groups_type(&DataType::Union( - UnionFields::try_new( - vec![0, 1], - vec![ - Field::new("a", DataType::Int32, true), - Field::new("b", DataType::Utf8, true), - ], - ) - .unwrap(), - UnionMode::Dense, - ))); - - let int_array = Int32Array::from(vec![5, 10, 5]); - let str_array = StringArray::from(vec!["foo", "foo"]); - let type_ids = vec![0_i8, 1, 0, 0, 1].into(); - let offsets = vec![0, 0, 1, 2, 1].into(); - let children = vec![ - Arc::new(int_array) as ArrayRef, - Arc::new(str_array) as ArrayRef, - ]; - let union_fields = [ - (0, Arc::new(Field::new("a", DataType::Int32, true))), - (1, Arc::new(Field::new("b", DataType::Utf8, true))), - ] - .into_iter() - .collect(); - let array: ArrayRef = Arc::new( - UnionArray::try_new(union_fields, type_ids, Some(offsets), children) - .unwrap(), - ); - - let mut acc = HLLAccumulator::new(); - acc.update_batch(&[Arc::clone(&array)]).unwrap(); - let non_group_count = distinct_count(&mut acc); - - let group_indices = vec![0usize; array.len()]; - let mut group_acc = HllGroupsAccumulator::new(); - group_acc - .update_batch(std::slice::from_ref(&array), &group_indices, None, 1) - .unwrap(); - let group_count = group_acc - .evaluate(EmitTo::All) - .unwrap() - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - - assert_eq!( - non_group_count, group_count, - "non-grouped and grouped paths disagree for union type" - ); - assert_eq!(non_group_count, 3, "wrong distinct count for union type"); - } - #[test] fn interval_support_numerical_acc_and_group_acc() { let year_month: ArrayRef =