This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 1c8295c992 feat: Support Union type in approx_distinct (#23714)
1c8295c992 is described below
commit 1c8295c992c597a53dcfa01f6dcf1cdc51268adb
Author: Michael Kleen <[email protected]>
AuthorDate: Thu Jul 23 10:29:37 2026 +0200
feat: Support Union type in approx_distinct (#23714)
## Which issue does this PR close?
- Relates to https://github.com/apache/datafusion/issues/22989 but does
not close it. More types are coming.
## Rationale for this change
- Support the Arrow type `Union` for `approx_distinct`
## What changes are included in this PR?
- Enable `HLLAccumulator` and `HllGroupsAccumulator` to support `Union`
- Tests
## Are these changes tested?
Yes
## Are there any user-facing changes?
Yes, `approx_distinct` supports now `Union` but no breaking changes.
---
.../functions-aggregate/src/approx_distinct.rs | 2 ++
datafusion/sqllogictest/src/test_context.rs | 41 ++++++++++++++++++++++
datafusion/sqllogictest/test_files/aggregate.slt | 30 ++++++++++++++++
3 files changed, 73 insertions(+)
diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs
b/datafusion/functions-aggregate/src/approx_distinct.rs
index 6d3ea3a8dd..f36c658e7f 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 e0aaa91ef6..fdb04edc05 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 cbb9c5d031..365505e2da 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<i: Int32, s: Utf8>) 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;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
