Skip to content

feat: Support Union type in approx_distinct #23714

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mkleen wants to merge 2 commits into
base: main
Choose a base branch
from
+62 −0
Open

feat: Support Union type in approx_distinct #23714

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ba3810b
feat: Support Union type in approx_distinct
mkleen Jul 19, 2026
9096c52
Merge branch 'main' into union_approx_distinct
mkleen Jul 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 62 additions & 0 deletions datafusion/functions-aggregate/src/approx_distinct.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))))
Expand Down Expand Up @@ -922,6 +923,7 @@ fn is_hll_groups_type(data_type: &DataType) -> bool {
| DataType::LargeListView(_)
| DataType::Map(_, _)
| DataType::Struct(_)
| DataType::Union(_, _)
)
}

Expand Down Expand Up @@ -1078,6 +1080,66 @@ mod tests {
assert_count_numerical_acc_and_group_acc::<Decimal256Type>(decimal_256, 6);
}

#[test]
fn union_support_hll_acc_and_group_acc() {

@Jefffrey Jefffrey Jul 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not possible to do via SLT?

@mkleen mkleen Jul 20, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's tricky via SLT. There is no SQL union type but only union operations in datafusion.

@Jefffrey Jefffrey Jul 21, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what if we add more rows here:

datafusion/datafusion/sqllogictest/src/test_context.rs

Lines 560 to 594 in 3a2bc0a

fn register_union_table(ctx: &SessionContext) {
let union = UnionArray::try_new(
UnionFields::try_new(
// typeids: 3 for int, 1 for string
vec![3, 1],
vec![
Field::new("int", DataType::Int32, false),
Field::new("string", DataType::Utf8, false),
],
)
.unwrap(),
ScalarBuffer::from(vec![3, 1, 3]),
None,
vec![
Arc::new(Int32Array::from(vec![1, 2, 3])),
Arc::new(StringArray::from(vec![
Some("foo"),
Some("bar"),
Some("baz"),
])),
],
)
.unwrap();
let schema = Schema::new(vec![Field::new(
"union_column",
union.data_type().clone(),
false,
)]);
let batch =
RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(union)]).unwrap();
ctx.register_batch("union_table", batch).unwrap();
}

then can use in slts like here

datafusion/datafusion/sqllogictest/test_files/union_function.slt

Lines 18 to 30 in 3a2bc0a

# Note: union_table is registered via Rust code in the sqllogictest test harness
# because there is no way to create a union type in SQL today
##########
## UNION DataType Tests
##########
query ?I
select union_column, union_extract(union_column, 'int') from union_table;
----
{int=1} 1
{string=bar} NULL
{int=3} 3

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::<UInt64Array>()
.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 =
Expand Down
Loading