feat: Support Union type in approx_distinct

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(_, _)
     )
 }
 

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;

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;