diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 6d3ea3a8ddded..f36c658e7f385 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(_, _) ) } 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;