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 =