1use crate::hyperloglog::{HLL_HASH_STATE, HyperLogLog, NUM_REGISTERS, count_from_hashes};
21use arrow::array::{
22 Array, ArrayRef, BinaryArray, BinaryBuilder, BooleanArray, PrimitiveArray,
23 UInt64Array,
24};
25use arrow::buffer::NullBuffer;
26use arrow::datatypes::{
27 ArrowPrimitiveType, DataType, Date32Type, Date64Type, Decimal32Type, Decimal64Type,
28 Decimal128Type, Decimal256Type, DurationMicrosecondType, DurationMillisecondType,
29 DurationNanosecondType, DurationSecondType, Field, FieldRef, Int32Type, Int64Type,
30 IntervalDayTimeType, IntervalMonthDayNanoType, IntervalUnit, IntervalYearMonthType,
31 Time32MillisecondType, Time32SecondType, Time64MicrosecondType, Time64NanosecondType,
32 TimeUnit, TimestampMicrosecondType, TimestampMillisecondType,
33 TimestampNanosecondType, TimestampSecondType, UInt32Type, UInt64Type,
34};
35use datafusion_common::ScalarValue;
36use datafusion_common::hash_utils::create_hashes;
37use datafusion_common::{
38 DataFusionError, Result, downcast_value, internal_datafusion_err, internal_err,
39 not_impl_err,
40};
41use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
42use datafusion_expr::utils::format_state_name;
43use datafusion_expr::{
44 Accumulator, AggregateUDFImpl, Documentation, EmitTo, GroupsAccumulator, Signature,
45 Volatility,
46};
47use datafusion_functions_aggregate_common::aggregate::count_distinct::{
48 Bitmap65536DistinctCountAccumulator, Bitmap65536DistinctCountAccumulatorI16,
49 BoolArray256DistinctCountAccumulator, BoolArray256DistinctCountAccumulatorI8,
50 BooleanDistinctCountAccumulator,
51};
52use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::filter_to_nulls;
53use datafusion_functions_aggregate_common::noop_accumulator::NoopAccumulator;
54use datafusion_macros::user_doc;
55use std::fmt::{Debug, Formatter};
56use std::hash::Hash;
57use std::mem::{size_of, size_of_val};
58use std::sync::Arc;
59
60make_udaf_expr_and_func!(
61 ApproxDistinct,
62 approx_distinct,
63 expression,
64 "approximate number of distinct input values",
65 approx_distinct_udaf
66);
67
68impl<T: Hash + ?Sized> From<&HyperLogLog<T>> for ScalarValue {
69 fn from(v: &HyperLogLog<T>) -> ScalarValue {
70 let values = v.as_ref().to_vec();
71 ScalarValue::Binary(Some(values))
72 }
73}
74
75impl<T: Hash + ?Sized> TryFrom<&[u8]> for HyperLogLog<T> {
76 type Error = DataFusionError;
77 fn try_from(v: &[u8]) -> Result<HyperLogLog<T>> {
78 let arr: [u8; 16384] = v.try_into().map_err(|_| {
79 internal_datafusion_err!("Impossibly got invalid binary array from states")
80 })?;
81 Ok(HyperLogLog::<T>::new_with_registers(arr))
82 }
83}
84
85impl<T: Hash + ?Sized> TryFrom<&ScalarValue> for HyperLogLog<T> {
86 type Error = DataFusionError;
87 fn try_from(v: &ScalarValue) -> Result<HyperLogLog<T>> {
88 if let ScalarValue::Binary(Some(slice)) = v {
89 slice.as_slice().try_into()
90 } else {
91 internal_err!(
92 "Impossibly got invalid scalar value while converting to HyperLogLog"
93 )
94 }
95 }
96}
97
98#[derive(Debug)]
99struct ApproxDistinctBitmapWrapper<A: Accumulator> {
100 inner: A,
101}
102
103impl<A: Accumulator> Accumulator for ApproxDistinctBitmapWrapper<A> {
104 fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
105 self.inner.update_batch(values)
106 }
107
108 fn evaluate(&mut self) -> Result<ScalarValue> {
109 match self.inner.evaluate()? {
110 ScalarValue::Int64(Some(v)) => Ok(ScalarValue::UInt64(Some(v as u64))),
111 other => internal_err!("unexpected: {other}"),
112 }
113 }
114
115 fn size(&self) -> usize {
116 self.inner.size()
117 }
118
119 fn state(&mut self) -> Result<Vec<ScalarValue>> {
120 self.inner.state()
121 }
122
123 fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
124 self.inner.merge_batch(states)
125 }
126}
127
128#[derive(Debug)]
129struct HLLAccumulator {
130 hll: HyperLogLog<u8>,
131 hashes: Vec<u64>,
132}
133
134impl HLLAccumulator {
135 pub fn new() -> Self {
136 Self {
137 hll: HyperLogLog::new(),
138 hashes: Vec::new(),
139 }
140 }
141}
142
143impl Accumulator for HLLAccumulator {
144 fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
145 let array = values[0].as_ref();
146 self.hashes.clear();
147 self.hashes.resize(array.len(), 0);
148 create_hashes([array], &HLL_HASH_STATE, &mut self.hashes)?;
149
150 match array.logical_nulls() {
151 None => {
152 for &hash in &self.hashes {
153 self.hll.add_hashed(hash);
154 }
155 }
156 Some(nulls) => {
157 for row in 0..array.len() {
158 if nulls.is_valid(row) {
159 self.hll.add_hashed(self.hashes[row]);
160 }
161 }
162 }
163 }
164 Ok(())
165 }
166
167 fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
168 assert_eq!(1, states.len(), "expect only 1 element in the states");
169 let binary_array = downcast_value!(states[0], BinaryArray);
170 for v in binary_array.iter() {
171 let v = v.ok_or_else(|| {
172 internal_datafusion_err!("Impossibly got empty binary array from states")
173 })?;
174 let other = v.try_into()?;
175 self.hll.merge(&other);
176 }
177 Ok(())
178 }
179
180 fn state(&mut self) -> Result<Vec<ScalarValue>> {
181 let value = ScalarValue::from(&self.hll);
182 Ok(vec![value])
183 }
184
185 fn evaluate(&mut self) -> Result<ScalarValue> {
186 Ok(ScalarValue::UInt64(Some(self.hll.count() as u64)))
187 }
188
189 fn size(&self) -> usize {
190 size_of_val(self) + self.hashes.capacity() * size_of::<u64>()
191 }
192}
193
194#[derive(Debug)]
196struct NumericHLLAccumulator<T>
197where
198 T: ArrowPrimitiveType,
199 T::Native: Hash,
200{
201 hll: HyperLogLog<T::Native>,
202}
203
204impl<T> NumericHLLAccumulator<T>
205where
206 T: ArrowPrimitiveType,
207 T::Native: Hash,
208{
209 pub fn new() -> Self {
210 Self {
211 hll: HyperLogLog::new(),
212 }
213 }
214}
215
216impl<T> Accumulator for NumericHLLAccumulator<T>
217where
218 T: ArrowPrimitiveType + Debug,
219 T::Native: Hash,
220{
221 fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
222 let array: &PrimitiveArray<T> = downcast_value!(values[0], PrimitiveArray, T);
223 self.hll.extend(array.into_iter().flatten());
224 Ok(())
225 }
226
227 fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
228 assert_eq!(1, states.len(), "expect only 1 element in the states");
229 let binary_array = downcast_value!(states[0], BinaryArray);
230 for v in binary_array.iter() {
231 let v = v.ok_or_else(|| {
232 internal_datafusion_err!("Impossibly got empty binary array from states")
233 })?;
234 let other = v.try_into()?;
235 self.hll.merge(&other);
236 }
237 Ok(())
238 }
239
240 fn state(&mut self) -> Result<Vec<ScalarValue>> {
241 let value = ScalarValue::from(&self.hll);
242 Ok(vec![value])
243 }
244
245 fn evaluate(&mut self) -> Result<ScalarValue> {
246 Ok(ScalarValue::UInt64(Some(self.hll.count() as u64)))
247 }
248
249 fn size(&self) -> usize {
250 size_of_val(self)
251 }
252}
253
254const SPARSE_LIMIT: usize = 256;
263
264#[derive(Clone, Debug)]
273enum GroupHll {
274 Sparse(Vec<u64>),
276 Dense(Box<HyperLogLog<u8>>),
277}
278
279impl Default for GroupHll {
280 fn default() -> Self {
281 GroupHll::Sparse(Vec::new())
282 }
283}
284
285fn fold_sparse_to_hll(hashes: &[u64]) -> HyperLogLog<u8> {
287 let mut hll = HyperLogLog::<u8>::new();
288 for &h in hashes {
289 hll.add_hashed(h);
290 }
291 hll
292}
293
294impl GroupHll {
295 #[inline]
298 fn add_hash(&mut self, hash: u64) -> isize {
299 match self {
300 GroupHll::Dense(hll) => {
301 hll.add_hashed(hash);
302 0
303 }
304 GroupHll::Sparse(v) => {
305 let cap_before = v.capacity();
306 v.push(hash);
307 if v.len() >= 2 * SPARSE_LIMIT {
308 return self.compact_or_promote(cap_before);
309 }
310 ((v.capacity() - cap_before) * size_of::<u64>()) as isize
311 }
312 }
313 }
314
315 #[cold]
318 fn compact_or_promote(&mut self, cap_before: usize) -> isize {
319 let GroupHll::Sparse(v) = self else {
320 return 0;
321 };
322 v.sort_unstable();
323 v.dedup();
324 if v.len() > SPARSE_LIMIT {
325 *self = GroupHll::Dense(Box::new(fold_sparse_to_hll(v)));
329 (NUM_REGISTERS as isize) - ((cap_before * size_of::<u64>()) as isize)
330 } else {
331 ((v.capacity() - cap_before) * size_of::<u64>()) as isize
334 }
335 }
336
337 fn merge_serialized(&mut self, bytes: &[u8]) -> Result<isize> {
340 if bytes.is_empty() {
341 return Ok(0);
342 }
343 if bytes.len() == NUM_REGISTERS {
344 let other: HyperLogLog<u8> = bytes.try_into()?;
345 Ok(self.merge_dense(&other))
346 } else {
347 if !bytes.len().is_multiple_of(size_of::<u64>()) {
348 return internal_err!(
349 "approx_distinct: malformed sparse state: length {} is not a multiple of {}",
350 bytes.len(),
351 size_of::<u64>()
352 );
353 }
354 if bytes.len() > SPARSE_LIMIT * size_of::<u64>() {
355 return internal_err!(
356 "approx_distinct: malformed sparse state: length {} exceeds sparse limit {}",
357 bytes.len(),
358 SPARSE_LIMIT * size_of::<u64>()
359 );
360 }
361 let mut delta = 0;
362 for chunk in bytes.chunks_exact(size_of::<u64>()) {
363 let h = u64::from_le_bytes(chunk.try_into().unwrap());
364 delta += self.add_hash(h);
365 }
366 Ok(delta)
367 }
368 }
369
370 fn merge_dense(&mut self, other: &HyperLogLog<u8>) -> isize {
372 match self {
373 GroupHll::Dense(hll) => {
374 hll.merge(other);
375 0
376 }
377 GroupHll::Sparse(v) => {
378 let cap_before = v.capacity();
379 let mut hll = other.clone();
380 for &h in v.iter() {
381 hll.add_hashed(h);
382 }
383 *self = GroupHll::Dense(Box::new(hll));
384 (NUM_REGISTERS as isize) - ((cap_before * size_of::<u64>()) as isize)
385 }
386 }
387 }
388
389 fn count(&self) -> u64 {
391 match self {
392 GroupHll::Dense(hll) => hll.count() as u64,
393 GroupHll::Sparse(v) => count_from_hashes(v) as u64,
397 }
398 }
399
400 fn heap_bytes(&self) -> usize {
404 match self {
405 GroupHll::Sparse(v) => v.capacity() * size_of::<u64>(),
406 GroupHll::Dense(_) => NUM_REGISTERS,
407 }
408 }
409
410 fn serialize(&mut self, scratch: &mut Vec<u8>) {
417 scratch.clear();
418 match self {
419 GroupHll::Dense(hll) => {
420 let registers: &[u8] = (**hll).as_ref();
421 scratch.extend_from_slice(registers);
422 }
423 GroupHll::Sparse(v) => {
424 v.sort_unstable();
425 v.dedup();
426 if v.len() > SPARSE_LIMIT {
427 scratch.extend_from_slice(fold_sparse_to_hll(v).as_ref());
428 } else {
429 for &h in v.iter() {
430 scratch.extend_from_slice(&h.to_le_bytes());
431 }
432 }
433 }
434 }
435 }
436}
437
438struct HllGroupsAccumulator {
470 groups: Vec<GroupHll>,
472 allocated_bytes: usize,
474 hashes: Vec<u64>,
476}
477
478impl HllGroupsAccumulator {
479 fn new() -> Self {
480 Self {
481 groups: Vec::new(),
482 allocated_bytes: 0,
483 hashes: Vec::new(),
484 }
485 }
486
487 #[inline]
488 fn ensure_groups(&mut self, total_num_groups: usize) {
489 if total_num_groups > self.groups.len() {
490 self.groups.resize_with(total_num_groups, GroupHll::default);
491 }
492 }
493
494 #[inline]
495 fn apply_delta(&mut self, delta: isize) {
496 self.allocated_bytes =
497 (self.allocated_bytes as isize).saturating_add(delta).max(0) as usize;
498 }
499}
500
501impl GroupsAccumulator for HllGroupsAccumulator {
502 fn update_batch(
503 &mut self,
504 values: &[ArrayRef],
505 group_indices: &[usize],
506 opt_filter: Option<&BooleanArray>,
507 total_num_groups: usize,
508 ) -> Result<()> {
509 self.ensure_groups(total_num_groups);
510 let array = values[0].as_ref();
511 self.hashes.clear();
512 self.hashes.resize(array.len(), 0);
513 create_hashes([array], &HLL_HASH_STATE, &mut self.hashes)?;
514
515 let mut delta: isize = 0;
516 let filter_nulls = opt_filter.map(filter_to_nulls);
519 let value_nulls = array.logical_nulls();
520 let combined_nulls =
521 NullBuffer::union(filter_nulls.as_ref(), value_nulls.as_ref());
522 match combined_nulls {
523 None => {
524 for (row, &hash) in self.hashes.iter().enumerate() {
525 delta += self.groups[group_indices[row]].add_hash(hash);
526 }
527 }
528 Some(nulls) => {
529 for row in nulls.valid_indices() {
530 delta += self.groups[group_indices[row]].add_hash(self.hashes[row]);
531 }
532 }
533 }
534 self.apply_delta(delta);
535 Ok(())
536 }
537
538 fn merge_batch(
539 &mut self,
540 values: &[ArrayRef],
541 group_indices: &[usize],
542 total_num_groups: usize,
543 ) -> Result<()> {
544 self.ensure_groups(total_num_groups);
545 let states = downcast_value!(values[0], BinaryArray);
546 let mut delta: isize = 0;
547 for (row, &group_index) in group_indices.iter().enumerate() {
548 if states.is_valid(row) {
549 delta += self.groups[group_index].merge_serialized(states.value(row))?;
550 }
551 }
552 self.apply_delta(delta);
553 Ok(())
554 }
555
556 fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
557 let groups = emit_to.take_needed(&mut self.groups);
558 let mut freed = 0;
559 let counts: UInt64Array = groups
560 .iter()
561 .map(|g| {
562 freed += g.heap_bytes();
563 Some(g.count())
564 })
565 .collect();
566 self.allocated_bytes = self.allocated_bytes.saturating_sub(freed);
568 Ok(Arc::new(counts))
569 }
570
571 fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
572 let mut groups = emit_to.take_needed(&mut self.groups);
573 let mut builder = BinaryBuilder::new();
574 let mut scratch: Vec<u8> = Vec::new();
575 let mut freed = 0;
576 for g in groups.iter_mut() {
577 freed += g.heap_bytes();
578 g.serialize(&mut scratch);
579 builder.append_value(&scratch);
580 }
581 self.allocated_bytes = self.allocated_bytes.saturating_sub(freed);
583 Ok(vec![Arc::new(builder.finish())])
584 }
585
586 fn convert_to_state(
587 &self,
588 values: &[ArrayRef],
589 opt_filter: Option<&BooleanArray>,
590 ) -> Result<Vec<ArrayRef>> {
591 assert_eq!(values.len(), 1, "single argument to convert_to_state");
592 let array = values[0].as_ref();
593 let mut hashes = vec![0; array.len()];
594 create_hashes([array], &HLL_HASH_STATE, &mut hashes)?;
595
596 let filter_nulls = opt_filter.map(filter_to_nulls);
597 let value_nulls = array.logical_nulls();
598 let combined_nulls =
599 NullBuffer::union(filter_nulls.as_ref(), value_nulls.as_ref());
600
601 let mut builder = BinaryBuilder::new();
602 let mut scratch = Vec::new();
603 for (row, hash) in hashes.into_iter().enumerate() {
604 if combined_nulls
605 .as_ref()
606 .is_none_or(|nulls| nulls.is_valid(row))
607 {
608 scratch.clear();
609 scratch.extend_from_slice(&hash.to_le_bytes());
610 builder.append_value(&scratch);
611 } else {
612 builder.append_value([]);
613 }
614 }
615
616 Ok(vec![Arc::new(builder.finish())])
617 }
618 fn size(&self) -> usize {
619 self.groups.capacity() * size_of::<GroupHll>()
620 + self.allocated_bytes
621 + self.hashes.capacity() * size_of::<u64>()
622 }
623}
624
625impl Debug for ApproxDistinct {
626 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
627 f.debug_struct("ApproxDistinct")
628 .field("name", &self.name())
629 .field("signature", &self.signature)
630 .finish()
631 }
632}
633
634impl Default for ApproxDistinct {
635 fn default() -> Self {
636 Self::new()
637 }
638}
639
640#[user_doc(
641 doc_section(label = "Approximate Functions"),
642 description = "Returns the approximate number of distinct input values calculated using the HyperLogLog algorithm.",
643 syntax_example = "approx_distinct(expression)",
644 sql_example = r#"```sql
645> SELECT approx_distinct(column_name) FROM table_name;
646+-----------------------------------+
647| approx_distinct(column_name) |
648+-----------------------------------+
649| 42 |
650+-----------------------------------+
651```"#,
652 standard_argument(name = "expression",)
653)]
654#[derive(PartialEq, Eq, Hash)]
655pub struct ApproxDistinct {
656 signature: Signature,
657}
658
659impl ApproxDistinct {
660 pub fn new() -> Self {
661 Self {
662 signature: Signature::any(1, Volatility::Immutable),
663 }
664 }
665}
666
667#[cold]
668fn get_fixed_domain_approx_accumulator(
669 data_type: &DataType,
670) -> Result<Box<dyn Accumulator>> {
671 match data_type {
672 DataType::Boolean => Ok(Box::new(ApproxDistinctBitmapWrapper {
673 inner: BooleanDistinctCountAccumulator::new(),
674 })),
675 DataType::UInt8 => Ok(Box::new(ApproxDistinctBitmapWrapper {
676 inner: BoolArray256DistinctCountAccumulator::new(),
677 })),
678 DataType::Int8 => Ok(Box::new(ApproxDistinctBitmapWrapper {
679 inner: BoolArray256DistinctCountAccumulatorI8::new(),
680 })),
681 DataType::UInt16 => Ok(Box::new(ApproxDistinctBitmapWrapper {
682 inner: Bitmap65536DistinctCountAccumulator::new(),
683 })),
684 DataType::Int16 => Ok(Box::new(ApproxDistinctBitmapWrapper {
685 inner: Bitmap65536DistinctCountAccumulatorI16::new(),
686 })),
687 _ => internal_err!("unsupported small int type: {}", data_type),
688 }
689}
690
691#[cold]
692fn get_fixed_domain_state_field(
693 name: &str,
694 data_type: &DataType,
695) -> Result<Vec<FieldRef>> {
696 Ok(vec![
697 Field::new_list(
698 format_state_name(name, "approx_distinct"),
699 Field::new_list_field(data_type.clone(), true),
700 false,
701 )
702 .into(),
703 ])
704}
705
706impl AggregateUDFImpl for ApproxDistinct {
707 fn name(&self) -> &str {
708 "approx_distinct"
709 }
710
711 fn signature(&self) -> &Signature {
712 &self.signature
713 }
714
715 fn return_type(&self, _: &[DataType]) -> Result<DataType> {
716 Ok(DataType::UInt64)
717 }
718
719 fn default_value(&self, _data_type: &DataType) -> Result<ScalarValue> {
720 Ok(ScalarValue::UInt64(Some(0)))
721 }
722
723 fn is_nullable(&self) -> bool {
724 false
725 }
726
727 fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
728 let data_type = args.input_fields[0].data_type();
729 match data_type {
730 DataType::Null => Ok(vec![
731 Field::new(
732 format_state_name(args.name, self.name()),
733 DataType::Null,
734 true,
735 )
736 .into(),
737 ]),
738 DataType::Boolean
739 | DataType::UInt8
740 | DataType::Int8
741 | DataType::UInt16
742 | DataType::Int16 => get_fixed_domain_state_field(args.name, data_type),
743 _ => Ok(vec![
744 Field::new(
745 format_state_name(args.name, "hll_registers"),
746 DataType::Binary,
747 false,
748 )
749 .into(),
750 ]),
751 }
752 }
753
754 fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
755 let data_type = acc_args.expr_fields[0].data_type();
756
757 let accumulator: Box<dyn Accumulator> = match data_type {
759 DataType::Boolean
760 | DataType::UInt8
761 | DataType::Int8
762 | DataType::UInt16
763 | DataType::Int16 => {
764 return get_fixed_domain_approx_accumulator(data_type);
765 }
766 DataType::UInt32 => Box::new(NumericHLLAccumulator::<UInt32Type>::new()),
767 DataType::UInt64 => Box::new(NumericHLLAccumulator::<UInt64Type>::new()),
768 DataType::Int32 => Box::new(NumericHLLAccumulator::<Int32Type>::new()),
769 DataType::Int64 => Box::new(NumericHLLAccumulator::<Int64Type>::new()),
770 DataType::Date32 => Box::new(NumericHLLAccumulator::<Date32Type>::new()),
771 DataType::Date64 => Box::new(NumericHLLAccumulator::<Date64Type>::new()),
772 DataType::Time32(TimeUnit::Second) => {
773 Box::new(NumericHLLAccumulator::<Time32SecondType>::new())
774 }
775 DataType::Time32(TimeUnit::Millisecond) => {
776 Box::new(NumericHLLAccumulator::<Time32MillisecondType>::new())
777 }
778 DataType::Time64(TimeUnit::Microsecond) => {
779 Box::new(NumericHLLAccumulator::<Time64MicrosecondType>::new())
780 }
781 DataType::Time64(TimeUnit::Nanosecond) => {
782 Box::new(NumericHLLAccumulator::<Time64NanosecondType>::new())
783 }
784 DataType::Timestamp(TimeUnit::Second, _) => {
785 Box::new(NumericHLLAccumulator::<TimestampSecondType>::new())
786 }
787 DataType::Timestamp(TimeUnit::Millisecond, _) => {
788 Box::new(NumericHLLAccumulator::<TimestampMillisecondType>::new())
789 }
790 DataType::Timestamp(TimeUnit::Microsecond, _) => {
791 Box::new(NumericHLLAccumulator::<TimestampMicrosecondType>::new())
792 }
793 DataType::Timestamp(TimeUnit::Nanosecond, _) => {
794 Box::new(NumericHLLAccumulator::<TimestampNanosecondType>::new())
795 }
796 DataType::Interval(IntervalUnit::YearMonth) => {
797 Box::new(NumericHLLAccumulator::<IntervalYearMonthType>::new())
798 }
799 DataType::Interval(IntervalUnit::DayTime) => {
800 Box::new(NumericHLLAccumulator::<IntervalDayTimeType>::new())
801 }
802 DataType::Interval(IntervalUnit::MonthDayNano) => {
803 Box::new(NumericHLLAccumulator::<IntervalMonthDayNanoType>::new())
804 }
805 DataType::Decimal32(_, _) => {
806 Box::new(NumericHLLAccumulator::<Decimal32Type>::new())
807 }
808 DataType::Decimal64(_, _) => {
809 Box::new(NumericHLLAccumulator::<Decimal64Type>::new())
810 }
811 DataType::Decimal128(_, _) => {
812 Box::new(NumericHLLAccumulator::<Decimal128Type>::new())
813 }
814 DataType::Decimal256(_, _) => {
815 Box::new(NumericHLLAccumulator::<Decimal256Type>::new())
816 }
817 DataType::Duration(TimeUnit::Second) => {
818 Box::new(NumericHLLAccumulator::<DurationSecondType>::new())
819 }
820 DataType::Duration(TimeUnit::Millisecond) => {
821 Box::new(NumericHLLAccumulator::<DurationMillisecondType>::new())
822 }
823 DataType::Duration(TimeUnit::Microsecond) => {
824 Box::new(NumericHLLAccumulator::<DurationMicrosecondType>::new())
825 }
826 DataType::Duration(TimeUnit::Nanosecond) => {
827 Box::new(NumericHLLAccumulator::<DurationNanosecondType>::new())
828 }
829 DataType::Utf8
830 | DataType::LargeUtf8
831 | DataType::Utf8View
832 | DataType::Binary
833 | DataType::BinaryView
834 | DataType::FixedSizeBinary(_)
835 | DataType::List(_)
836 | DataType::LargeList(_)
837 | DataType::FixedSizeList(_, _)
838 | DataType::ListView(_)
839 | DataType::LargeListView(_)
840 | DataType::Map(_, _)
841 | DataType::Struct(_)
842 | DataType::Union(_, _)
843 | DataType::LargeBinary => Box::new(HLLAccumulator::new()),
844 DataType::Null => {
845 Box::new(NoopAccumulator::new(ScalarValue::UInt64(Some(0))))
846 }
847 other => {
848 return not_impl_err!(
849 "Support for 'approx_distinct' for data type {other} is not implemented"
850 );
851 }
852 };
853 Ok(accumulator)
854 }
855
856 fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
857 is_hll_groups_type(args.expr_fields[0].data_type())
858 }
859
860 fn create_groups_accumulator(
861 &self,
862 args: AccumulatorArgs,
863 ) -> Result<Box<dyn GroupsAccumulator>> {
864 let data_type = args.expr_fields[0].data_type();
865 if is_hll_groups_type(data_type) {
866 Ok(Box::new(HllGroupsAccumulator::new()))
867 } else {
868 not_impl_err!(
869 "GroupsAccumulator for 'approx_distinct' is not implemented for data type {data_type}"
870 )
871 }
872 }
873
874 fn documentation(&self) -> Option<&Documentation> {
875 self.doc()
876 }
877}
878
879fn is_hll_groups_type(data_type: &DataType) -> bool {
883 matches!(
884 data_type,
885 DataType::UInt32
886 | DataType::UInt64
887 | DataType::Int32
888 | DataType::Int64
889 | DataType::Date32
890 | DataType::Date64
891 | DataType::Time32(TimeUnit::Second)
892 | DataType::Time32(TimeUnit::Millisecond)
893 | DataType::Time64(TimeUnit::Microsecond)
894 | DataType::Time64(TimeUnit::Nanosecond)
895 | DataType::Timestamp(TimeUnit::Second, _)
896 | DataType::Timestamp(TimeUnit::Millisecond, _)
897 | DataType::Timestamp(TimeUnit::Microsecond, _)
898 | DataType::Timestamp(TimeUnit::Nanosecond, _)
899 | DataType::Interval(IntervalUnit::YearMonth)
900 | DataType::Interval(IntervalUnit::DayTime)
901 | DataType::Interval(IntervalUnit::MonthDayNano)
902 | DataType::Decimal32(_, _)
903 | DataType::Decimal64(_, _)
904 | DataType::Decimal128(_, _)
905 | DataType::Decimal256(_, _)
906 | DataType::Duration(_)
907 | DataType::Utf8
908 | DataType::LargeUtf8
909 | DataType::Utf8View
910 | DataType::Binary
911 | DataType::BinaryView
912 | DataType::FixedSizeBinary(_)
913 | DataType::LargeBinary
914 | DataType::List(_)
915 | DataType::LargeList(_)
916 | DataType::FixedSizeList(_, _)
917 | DataType::ListView(_)
918 | DataType::LargeListView(_)
919 | DataType::Map(_, _)
920 | DataType::Struct(_)
921 | DataType::Union(_, _)
922 )
923}
924
925#[cfg(test)]
926mod tests {
927 use super::*;
928 use std::hash::BuildHasher;
929
930 #[cfg(not(feature = "force_hash_collisions"))]
931 mod real_hash_test {
932 use super::*;
933 use arrow::array::{
934 AsArray, Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array,
935 Int64Array, IntervalDayTimeArray, IntervalMonthDayNanoArray,
936 IntervalYearMonthArray, StringViewArray,
937 };
938 use arrow::datatypes::{IntervalDayTime, IntervalMonthDayNano, i256};
939 use std::sync::Arc;
940 const LONG: &str = "this string is definitely longer than twelve bytes";
942
943 fn distinct_count(acc: &mut HLLAccumulator) -> u64 {
944 match acc.evaluate().unwrap() {
945 ScalarValue::UInt64(Some(v)) => v,
946 other => panic!("unexpected evaluate result: {other:?}"),
947 }
948 }
949
950 fn assert_count_numerical_acc_and_group_acc<T>(array: ArrayRef, expected: u64)
951 where
952 T: ArrowPrimitiveType + Debug,
953 T::Native: Hash,
954 {
955 assert!(
956 is_hll_groups_type(array.data_type()),
957 "{} should be groups-capable",
958 array.data_type()
959 );
960
961 let mut acc = NumericHLLAccumulator::<T>::new();
962 acc.update_batch(&[Arc::clone(&array)]).unwrap();
963 let per_group_count = match acc.evaluate().unwrap() {
964 ScalarValue::UInt64(Some(v)) => v,
965 other => panic!("unexpected evaluate result: {other:?}"),
966 };
967
968 let group_indices = vec![0usize; array.len()];
969 let mut acc = HllGroupsAccumulator::new();
970 acc.update_batch(std::slice::from_ref(&array), &group_indices, None, 1)
971 .unwrap();
972 let groups_count = acc
973 .evaluate(EmitTo::All)
974 .unwrap()
975 .as_any()
976 .downcast_ref::<UInt64Array>()
977 .unwrap()
978 .value(0);
979
980 assert_eq!(
981 per_group_count,
982 groups_count,
983 "paths disagree for {}",
984 array.data_type()
985 );
986 assert_eq!(
987 per_group_count,
988 expected,
989 "wrong count for {}",
990 array.data_type()
991 );
992 }
993
994 #[test]
995 fn decimal_support_numerical_acc_and_group_acc() {
996 let decimal_32: ArrayRef = Arc::new(
997 Decimal32Array::from(vec![
998 1i32,
999 2,
1000 2,
1001 3,
1002 3,
1003 3,
1004 0,
1005 0,
1006 123_456_789,
1007 999_999_999,
1008 999_999_999,
1009 ])
1010 .with_precision_and_scale(9, 2)
1011 .unwrap(),
1012 );
1013 assert_count_numerical_acc_and_group_acc::<Decimal32Type>(decimal_32, 6);
1014
1015 let decimal_64: ArrayRef = Arc::new(
1016 Decimal64Array::from(vec![
1017 1i64,
1018 2,
1019 2,
1020 3,
1021 3,
1022 3,
1023 0,
1024 0,
1025 1_234_567_890_123,
1026 9_999_999_999_999,
1027 9_999_999_999_999,
1028 ])
1029 .with_precision_and_scale(18, 2)
1030 .unwrap(),
1031 );
1032 assert_count_numerical_acc_and_group_acc::<Decimal64Type>(decimal_64, 6);
1033
1034 let decimal_128: ArrayRef = Arc::new(
1035 Decimal128Array::from(vec![
1036 1i128,
1037 2,
1038 2,
1039 3,
1040 3,
1041 3,
1042 0,
1043 0,
1044 1_234_567_890,
1045 9_999_999_999,
1046 9_999_999_999,
1047 ])
1048 .with_precision_and_scale(38, 2)
1049 .unwrap(),
1050 );
1051 assert_count_numerical_acc_and_group_acc::<Decimal128Type>(decimal_128, 6);
1052
1053 let big_256_a =
1054 i256::from_string("123456789012345678901234567890123456").unwrap();
1055 let big_256_b =
1056 i256::from_string("987654321098765432109876543210987654").unwrap();
1057
1058 let decimal_256: ArrayRef = Arc::new(
1059 Decimal256Array::from(vec![
1060 i256::from_i128(1),
1061 i256::from_i128(2),
1062 i256::from_i128(2),
1063 i256::from_i128(3),
1064 i256::from_i128(3),
1065 i256::from_i128(3),
1066 i256::from_i128(0),
1067 i256::from_i128(0),
1068 big_256_a,
1069 big_256_b,
1070 big_256_b,
1071 ])
1072 .with_precision_and_scale(40, 2)
1073 .unwrap(),
1074 );
1075 assert_count_numerical_acc_and_group_acc::<Decimal256Type>(decimal_256, 6);
1076 }
1077
1078 #[test]
1079 fn interval_support_numerical_acc_and_group_acc() {
1080 let year_month: ArrayRef =
1081 Arc::new(IntervalYearMonthArray::from(vec![1, 2, 2, 3, 3, 3, 0, 0]));
1082 assert_count_numerical_acc_and_group_acc::<IntervalYearMonthType>(
1083 year_month, 4,
1084 );
1085
1086 let day_time: ArrayRef = Arc::new(IntervalDayTimeArray::from(vec![
1087 IntervalDayTime::new(1, 0),
1088 IntervalDayTime::new(1, 0),
1089 IntervalDayTime::new(1, 5),
1090 IntervalDayTime::new(2, 0),
1091 ]));
1092 assert_count_numerical_acc_and_group_acc::<IntervalDayTimeType>(day_time, 3);
1093
1094 let month_day_nano: ArrayRef =
1095 Arc::new(IntervalMonthDayNanoArray::from(vec![
1096 IntervalMonthDayNano::new(1, 0, 0),
1097 IntervalMonthDayNano::new(1, 0, 0),
1098 IntervalMonthDayNano::new(1, 0, 5),
1099 IntervalMonthDayNano::new(0, 2, 0),
1100 IntervalMonthDayNano::new(0, 0, 0),
1101 ]));
1102 assert_count_numerical_acc_and_group_acc::<IntervalMonthDayNanoType>(
1103 month_day_nano,
1104 4,
1105 );
1106 }
1107
1108 #[test]
1111 fn update_batch_nullable_filter_excludes_null_filter_rows() {
1112 let values: ArrayRef = Arc::new(Int64Array::from(vec![1i64, 2, 3, 4, 5]));
1113 let filter =
1116 BooleanArray::from(vec![Some(true), None, Some(false), None, Some(true)]);
1117
1118 let mut acc = HllGroupsAccumulator::new();
1119 let group_indices = vec![0usize; 5];
1121 acc.update_batch(&[values], &group_indices, Some(&filter), 1)
1122 .unwrap();
1123
1124 let result = acc.evaluate(EmitTo::All).unwrap();
1126 let counts = result.as_any().downcast_ref::<UInt64Array>().unwrap();
1127 let expected = reference_count(&[h(1), h(5)]);
1129 assert_eq!(counts.value(0), expected);
1130 }
1131
1132 #[test]
1133 fn groups_convert_to_state_roundtrips_through_merge() {
1134 let values: ArrayRef = Arc::new(Int64Array::from(vec![
1135 Some(1),
1136 Some(2),
1137 Some(2),
1138 None,
1139 Some(3),
1140 ]));
1141 let filter = BooleanArray::from(vec![
1142 Some(true),
1143 Some(true),
1144 Some(true),
1145 Some(true),
1146 None,
1147 ]);
1148 let group_indices = vec![0usize, 1, 0, 1, 0];
1149
1150 let mut direct = HllGroupsAccumulator::new();
1151 direct
1152 .update_batch(
1153 std::slice::from_ref(&values),
1154 &group_indices,
1155 Some(&filter),
1156 2,
1157 )
1158 .unwrap();
1159 let direct = direct
1160 .evaluate(EmitTo::All)
1161 .unwrap()
1162 .as_any()
1163 .downcast_ref::<UInt64Array>()
1164 .unwrap()
1165 .clone();
1166
1167 let converter = HllGroupsAccumulator::new();
1168 let state = converter
1169 .convert_to_state(std::slice::from_ref(&values), Some(&filter))
1170 .unwrap();
1171 assert_eq!(state[0].null_count(), 0);
1172 let mut merged = HllGroupsAccumulator::new();
1173 merged.merge_batch(&state, &group_indices, 2).unwrap();
1174 let merged = merged
1175 .evaluate(EmitTo::All)
1176 .unwrap()
1177 .as_any()
1178 .downcast_ref::<UInt64Array>()
1179 .unwrap()
1180 .clone();
1181
1182 assert_eq!(direct, merged);
1183 }
1184
1185 #[test]
1186 fn groups_convert_to_state_preserves_empty_and_filtered_rows() {
1187 let converter = HllGroupsAccumulator::new();
1188 let empty_values: ArrayRef =
1189 Arc::new(Int64Array::from(Vec::<Option<i64>>::new()));
1190 let state = converter
1191 .convert_to_state(std::slice::from_ref(&empty_values), None)
1192 .unwrap();
1193 assert_eq!(state[0].len(), 0);
1194 assert_eq!(state[0].null_count(), 0);
1195
1196 let values: ArrayRef =
1197 Arc::new(Int64Array::from(vec![Some(1), Some(2), None]));
1198 let filter = BooleanArray::from(vec![Some(false), None, Some(false)]);
1199 let group_indices = vec![0usize, 1, 0];
1200 let state = converter
1201 .convert_to_state(std::slice::from_ref(&values), Some(&filter))
1202 .unwrap();
1203 assert_eq!(state[0].len(), values.len());
1204 assert_eq!(state[0].null_count(), 0);
1205 let state = state[0].as_any().downcast_ref::<BinaryArray>().unwrap();
1206 for row in 0..state.len() {
1207 assert_eq!(state.value(row), b"");
1208 }
1209
1210 let mut merged = HllGroupsAccumulator::new();
1211 merged
1212 .merge_batch(&[Arc::new(state.clone())], &group_indices, 2)
1213 .unwrap();
1214 let result = merged
1215 .evaluate(EmitTo::All)
1216 .unwrap()
1217 .as_any()
1218 .downcast_ref::<UInt64Array>()
1219 .unwrap()
1220 .clone();
1221 assert_eq!(result, UInt64Array::from(vec![0, 0]));
1222 }
1223
1224 #[test]
1228 fn utf8view_groups_short_string_hashed_consistently_across_batches() {
1229 let batch1: ArrayRef = Arc::new(StringViewArray::from(vec!["aaa", "bbb"]));
1231 assert!(batch1.as_string_view().data_buffers().is_empty());
1232
1233 let batch2: ArrayRef = Arc::new(StringViewArray::from(vec!["aaa", LONG]));
1236 assert!(!batch2.as_string_view().data_buffers().is_empty());
1237
1238 let group_indices = vec![0usize, 0];
1239 let mut acc = HllGroupsAccumulator::new();
1240 acc.update_batch(&[batch1], &group_indices, None, 1)
1241 .unwrap();
1242 acc.update_batch(&[batch2], &group_indices, None, 1)
1243 .unwrap();
1244
1245 let result = acc.evaluate(EmitTo::All).unwrap();
1247 let counts = result.as_any().downcast_ref::<UInt64Array>().unwrap();
1248 assert_eq!(counts.value(0), 3);
1249 }
1250
1251 #[test]
1254 fn utf8view_acc_split_batches_match_single_mixed_batch() {
1255 let mixed: ArrayRef =
1257 Arc::new(StringViewArray::from(vec!["aaa", "bbb", LONG, "aaa"]));
1258 let mut acc_single = HLLAccumulator::new();
1259 acc_single.update_batch(&[mixed]).unwrap();
1260
1261 let inline_only: ArrayRef =
1264 Arc::new(StringViewArray::from(vec!["aaa", "bbb"]));
1265 let with_buffer: ArrayRef =
1266 Arc::new(StringViewArray::from(vec!["aaa", LONG]));
1267 assert!(inline_only.as_string_view().data_buffers().is_empty());
1268 assert!(!with_buffer.as_string_view().data_buffers().is_empty());
1269
1270 let mut acc_split = HLLAccumulator::new();
1271 acc_split.update_batch(&[inline_only]).unwrap();
1272 acc_split.update_batch(&[with_buffer]).unwrap();
1273
1274 assert_eq!(
1275 distinct_count(&mut acc_single),
1276 distinct_count(&mut acc_split)
1277 );
1278 assert_eq!(distinct_count(&mut acc_single), 3);
1279 }
1280 }
1281
1282 fn h(v: u64) -> u64 {
1283 HLL_HASH_STATE.hash_one(v)
1284 }
1285
1286 fn reference_count(hashes: &[u64]) -> u64 {
1289 let mut hll = HyperLogLog::<u8>::new();
1290 for &hash in hashes {
1291 hll.add_hashed(hash);
1292 }
1293 hll.count() as u64
1294 }
1295
1296 fn serialize(g: &mut GroupHll) -> Vec<u8> {
1297 let mut buf = Vec::new();
1298 g.serialize(&mut buf);
1299 buf
1300 }
1301
1302 #[test]
1303 fn sparse_stays_sparse_for_small_groups() {
1304 let mut g = GroupHll::default();
1305 let hashes: Vec<u64> = (0..50).map(h).collect();
1306 for &hash in &hashes {
1307 g.add_hash(hash);
1308 }
1309 for &hash in &hashes {
1311 g.add_hash(hash);
1312 }
1313 assert!(
1314 matches!(g, GroupHll::Sparse(_)),
1315 "small group must be sparse"
1316 );
1317 assert_eq!(g.count(), reference_count(&hashes));
1318 let serialized = serialize(&mut g);
1321 assert!(serialized.len() < NUM_REGISTERS);
1322 assert!(serialized.len() <= SPARSE_LIMIT * size_of::<u64>());
1323 }
1324
1325 #[test]
1326 fn promotes_to_dense_for_large_groups() {
1327 let mut g = GroupHll::default();
1328 let hashes: Vec<u64> = (0..(SPARSE_LIMIT as u64 * 4)).map(h).collect();
1329 for &hash in &hashes {
1330 g.add_hash(hash);
1331 }
1332 assert!(matches!(g, GroupHll::Dense(_)), "large group must be dense");
1333 assert_eq!(g.count(), reference_count(&hashes));
1334 }
1335
1336 #[test]
1337 fn serialize_then_merge_roundtrips() {
1338 for n in [0u64, 10, SPARSE_LIMIT as u64 * 4] {
1339 let hashes: Vec<u64> = (0..n).map(h).collect();
1340 let mut src = GroupHll::default();
1341 for &hash in &hashes {
1342 src.add_hash(hash);
1343 }
1344 let bytes = serialize(&mut src);
1345 let mut dst = GroupHll::default();
1346 dst.merge_serialized(&bytes).unwrap();
1347 assert_eq!(dst.count(), reference_count(&hashes), "n = {n}");
1348 }
1349 }
1350
1351 #[test]
1352 fn sparse_limit_group_serializes_as_mergeable_sparse_state() {
1353 let hashes: Vec<u64> = (0..SPARSE_LIMIT as u64).map(h).collect();
1354 let mut src = GroupHll::default();
1355 for &hash in &hashes {
1356 src.add_hash(hash);
1357 }
1358 assert!(matches!(src, GroupHll::Sparse(_)));
1359
1360 let bytes = serialize(&mut src);
1361 assert_eq!(bytes.len(), SPARSE_LIMIT * size_of::<u64>());
1362
1363 let mut dst = GroupHll::default();
1364 dst.merge_serialized(&bytes).unwrap();
1365 assert_eq!(dst.count(), reference_count(&hashes));
1366 }
1367
1368 #[test]
1369 fn medium_sparse_group_serializes_as_mergeable_dense_state() {
1370 let n = SPARSE_LIMIT as u64 + 44;
1371 let hashes: Vec<u64> = (0..n).map(h).collect();
1372 let mut src = GroupHll::default();
1373 for &hash in &hashes {
1374 src.add_hash(hash);
1375 }
1376 assert!(
1377 matches!(src, GroupHll::Sparse(_)),
1378 "group should not promote during update before the compaction threshold"
1379 );
1380
1381 let bytes = serialize(&mut src);
1382 assert_eq!(bytes.len(), NUM_REGISTERS);
1383
1384 let mut dst = GroupHll::default();
1385 dst.merge_serialized(&bytes).unwrap();
1386 assert_eq!(dst.count(), reference_count(&hashes));
1387 }
1388
1389 #[test]
1390 fn merge_combines_disjoint_groups() {
1391 let left: Vec<u64> = (0..100).map(h).collect();
1393 let right: Vec<u64> = (100..(SPARSE_LIMIT as u64 * 4)).map(h).collect();
1394 let all: Vec<u64> = left.iter().chain(right.iter()).copied().collect();
1395
1396 let mut a = GroupHll::default();
1397 for &hash in &left {
1398 a.add_hash(hash);
1399 }
1400 let mut b = GroupHll::default();
1401 for &hash in &right {
1402 b.add_hash(hash);
1403 }
1404 let b_bytes = serialize(&mut b);
1405 a.merge_serialized(&b_bytes).unwrap();
1406 assert_eq!(a.count(), reference_count(&all));
1407 }
1408
1409 #[test]
1410 fn empty_group_counts_zero() {
1411 let mut g = GroupHll::default();
1412 assert_eq!(g.count(), 0);
1413 let bytes = serialize(&mut g);
1414 assert!(bytes.is_empty());
1415 let mut dst = GroupHll::default();
1416 dst.merge_serialized(&bytes).unwrap();
1417 assert_eq!(dst.count(), 0);
1418 }
1419}