Skip to main content

datafusion_functions_aggregate/
approx_distinct.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Defines physical expressions that can evaluated at runtime during query execution
19
20use 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/// Specialize the numeric case for extra performance.
195#[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
254/// Maximum number of distinct hashes kept in the sparse representation of a
255/// per-group sketch before it is promoted to a dense [`HyperLogLog`].
256///
257/// A dense sketch always occupies [`NUM_REGISTERS`] (16 KiB) regardless of how
258/// many values it has seen. The vast majority of groups in a high-cardinality
259/// `GROUP BY` only observe a handful of distinct values, so keeping their state
260/// as a small list of hashes saves a huge amount of memory (both while
261/// aggregating and when serializing the partial state for the final phase).
262const SPARSE_LIMIT: usize = 256;
263
264/// Per-group HyperLogLog state used by [`HllGroupsAccumulator`].
265///
266/// Starts out as a compact list of the (deduplicated) hashes observed for the
267/// group and only switches to a full dense [`HyperLogLog`] once it has seen more
268/// than [`SPARSE_LIMIT`] distinct values. Folding the stored hashes into a dense
269/// sketch produces exactly the same registers as adding the original values one
270/// by one, so the cardinality estimate is identical to the per-group
271/// [`Accumulator`] path.
272#[derive(Clone, Debug)]
273enum GroupHll {
274    /// Distinct hashes seen so far. May contain duplicates between compactions.
275    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
285/// Fold a slice of pre-computed hashes into a fresh [`HyperLogLog`] sketch.
286fn 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    /// Add a pre-computed hash, returning the change in heap-allocated bytes so
296    /// the accumulator can track its memory usage incrementally.
297    #[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    /// Deduplicate the sparse hash list and, if it still exceeds
316    /// [`SPARSE_LIMIT`] distinct values, promote it to a dense sketch.
317    #[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            // cap_before is the capacity already reflected in allocated_bytes.
326            // Any reallocation caused by the triggering push was never counted and
327            // is also freed here, so the two cancel out.
328            *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            // Account for any Vec growth caused by the triggering push.
332            // sort/dedup do not reallocate, so v.capacity() is the post-push capacity.
333            ((v.capacity() - cap_before) * size_of::<u64>()) as isize
334        }
335    }
336
337    /// Merge a serialized state (produced by [`Self::serialize`] or by the
338    /// per-group [`Accumulator`]) into this sketch.
339    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    /// Merge a dense sketch into this one, promoting to dense if necessary.
371    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    /// The approximate number of distinct values seen by this group.
390    fn count(&self) -> u64 {
391        match self {
392            GroupHll::Dense(hll) => hll.count() as u64,
393            // Estimate directly from the stored hashes; this produces exactly the
394            // same value as folding them into a dense sketch but avoids
395            // allocating and scanning a 16 KiB register array for every group.
396            GroupHll::Sparse(v) => count_from_hashes(v) as u64,
397        }
398    }
399
400    /// Heap bytes held by this sketch. Mirrors the deltas accrued in
401    /// [`Self::add_hash`] / [`Self::merge_dense`] so emitting a group can
402    /// precisely reverse them.
403    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    /// Serialize the sketch into `scratch` (which is cleared first). A dense
411    /// sketch is written as its raw [`NUM_REGISTERS`] registers (wire-compatible
412    /// with the per-group [`Accumulator`]); a sparse sketch is written as its
413    /// distinct hashes in little-endian order unless it has crossed
414    /// [`SPARSE_LIMIT`], in which case it is emitted as dense state so the final
415    /// merge path accepts it.
416    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
438/// A [`GroupsAccumulator`] for `approx_distinct` that keeps one adaptive
439/// (sparse → dense) HyperLogLog sketch per group.
440///
441/// This is dramatically faster than the generic `GroupsAccumulatorAdapter`
442/// fallback for high-cardinality `GROUP BY`s: it processes the whole input in a
443/// single vectorized pass (no per-group `take`/slice and no dynamic dispatch),
444/// and the sparse representation avoids allocating a 16 KiB sketch for every
445/// group when most groups only see a few distinct values.
446///
447///
448/// # Example
449///
450/// For `SELECT k, approx_distinct(v) FROM t GROUP BY k`, each group owns one
451/// independent sketch:
452///
453/// ```text
454/// group   state
455/// a       Sparse([h1, h2, h3, h2])
456/// b       Dense(HLL registers)
457/// ...
458/// ```
459///
460/// Group `a` has fewer than [`SPARSE_LIMIT`] distinct hashes, so it stays in
461/// the sparse representation. Before emitting state or estimating the count, the
462/// hash list is sorted and deduplicated to `[h1, h2, h3]`, then those hashes are
463/// interpreted exactly as if they had been added to a dense [`HyperLogLog`].
464///
465/// Group `b` has crossed the sparse limit, so its hashes have already been
466/// replayed into a dense sketch. New values for `b` update the dense registers
467/// directly, and serialized state is the raw [`NUM_REGISTERS`]-byte register
468/// array.
469struct HllGroupsAccumulator {
470    /// Per-group sketches, indexed by `group_index`.
471    groups: Vec<GroupHll>,
472    /// Incrementally maintained estimate of heap bytes used by `groups`.
473    allocated_bytes: usize,
474    /// Reused workspace for vectorized value hashing.
475    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        // Pre-combine value-nulls and filter into one mask so the update loop
517        // only visits rows that should affect the sketch.
518        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        // The emitted groups have been removed; reclaim their tracked bytes.
567        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        // The emitted groups have been removed; reclaim their tracked bytes.
582        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        // For primitive types, use specialized accumulators for better performance.
758        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
879/// Returns true for the data types backed by the HyperLogLog
880/// [`HllGroupsAccumulator`]. The fixed-domain types (booleans / small ints) and
881/// `Null` fall back to the per-group [`Accumulator`] path.
882fn 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        // A string longer than the 12-byte inline limit
941        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        /// `approx_distinct(v) FILTER (WHERE nullable_bool)` — a NULL filter row
1109        /// must not be counted (null filter is treated the same as false).
1110        #[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            // row 0: filter=true, row 1: filter=NULL, row 2: filter=false,
1114            // row 3: filter=NULL, row 4: filter=true
1115            let filter =
1116                BooleanArray::from(vec![Some(true), None, Some(false), None, Some(true)]);
1117
1118            let mut acc = HllGroupsAccumulator::new();
1119            // put all rows in group 0
1120            let group_indices = vec![0usize; 5];
1121            acc.update_batch(&[values], &group_indices, Some(&filter), 1)
1122                .unwrap();
1123
1124            // Only rows 0 and 4 (values 1 and 5) should be counted.
1125            let result = acc.evaluate(EmitTo::All).unwrap();
1126            let counts = result.as_any().downcast_ref::<UInt64Array>().unwrap();
1127            // reference: hash 1 and 5 into a dense sketch
1128            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        /// Regression: a short (≤ 12-byte) Utf8View string must hash identically
1225        /// in an all-inline batch and in a mixed batch that also contains a long
1226        /// string (which forces a data buffer).
1227        #[test]
1228        fn utf8view_groups_short_string_hashed_consistently_across_batches() {
1229            // Batch 1: all-inline (no data buffers) — "aaa" is hashed as u128 view.
1230            let batch1: ArrayRef = Arc::new(StringViewArray::from(vec!["aaa", "bbb"]));
1231            assert!(batch1.as_string_view().data_buffers().is_empty());
1232
1233            // Batch 2: mixed — LONG forces a data buffer; "aaa" must still be
1234            // hashed as u128 view so it matches its appearance in batch 1.
1235            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            // True distinct values: {"aaa", "bbb", LONG} == 3.
1246            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        /// Regression: a short (≤ 12-byte) Utf8View string must hash identically
1252        /// regardless of which batch it appears in — all-inline or mixed.
1253        #[test]
1254        fn utf8view_acc_split_batches_match_single_mixed_batch() {
1255            // Multiset: {"aaa" x2, "bbb", LONG}, so 3 distinct values.
1256            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            // Same multiset, but split so "aaa" lands in both an all-inline batch
1262            // and a batch with a data buffer (forced by LONG).
1263            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    /// Reference count: fold the given distinct hashes straight into a dense
1287    /// HyperLogLog. The grouped sketch must agree with this exactly.
1288    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        // duplicates must not change the estimate or trigger promotion
1310        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        // sparse serialized state is far smaller than a dense 16 KiB sketch
1319        // and must not exceed the sparse limit contract enforced by merge_serialized
1320        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        // sparse + sparse, sparse + dense, dense + dense
1392        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}