datafusion_functions_aggregate_common/aggregate/count_distinct/
native.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//! Specialized implementation of `COUNT DISTINCT` for "Native" arrays such as
19//! [`Int64Array`] and [`Float64Array`]
20//!
21//! [`Int64Array`]: arrow::array::Int64Array
22//! [`Float64Array`]: arrow::array::Float64Array
23use std::collections::HashSet;
24use std::fmt::Debug;
25use std::hash::Hash;
26use std::mem::size_of_val;
27use std::sync::Arc;
28
29use arrow::array::Array;
30use arrow::array::ArrayRef;
31use arrow::array::BooleanArray;
32use arrow::array::PrimitiveArray;
33use arrow::array::types::ArrowPrimitiveType;
34use arrow::datatypes::DataType;
35use datafusion_common::hash_utils::RandomState;
36
37use datafusion_common::ScalarValue;
38use datafusion_common::cast::{as_boolean_array, as_list_array, as_primitive_array};
39use datafusion_common::utils::SingleRowListArrayBuilder;
40use datafusion_common::utils::memory::estimate_memory_size;
41use datafusion_expr_common::accumulator::Accumulator;
42
43use crate::utils::GenericDistinctBuffer;
44
45#[derive(Debug)]
46pub struct PrimitiveDistinctCountAccumulator<T>
47where
48    T: ArrowPrimitiveType + Send,
49    T::Native: Eq + Hash,
50{
51    values: HashSet<T::Native, RandomState>,
52    data_type: DataType,
53}
54
55impl<T> PrimitiveDistinctCountAccumulator<T>
56where
57    T: ArrowPrimitiveType + Send,
58    T::Native: Eq + Hash,
59{
60    pub fn new(data_type: &DataType) -> Self {
61        Self {
62            values: HashSet::default(),
63            data_type: data_type.clone(),
64        }
65    }
66}
67
68impl<T> Accumulator for PrimitiveDistinctCountAccumulator<T>
69where
70    T: ArrowPrimitiveType + Send + Debug,
71    T::Native: Eq + Hash,
72{
73    fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
74        let arr = Arc::new(
75            PrimitiveArray::<T>::from_iter_values(self.values.iter().cloned())
76                .with_data_type(self.data_type.clone()),
77        );
78        Ok(vec![
79            SingleRowListArrayBuilder::new(arr).build_list_scalar(),
80        ])
81    }
82
83    #[inline(never)]
84    fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
85        if values.is_empty() {
86            return Ok(());
87        }
88
89        let arr = as_primitive_array::<T>(&values[0])?;
90        if arr.null_count() == 0 {
91            // Fast path: no nulls, so skip the per-element validity check and
92            // insert directly from the values buffer (mirrors `merge_batch`).
93            self.values.extend(arr.values().iter().copied());
94        } else {
95            arr.iter().flatten().for_each(|value| {
96                self.values.insert(value);
97            });
98        }
99
100        Ok(())
101    }
102
103    fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> {
104        if states.is_empty() {
105            return Ok(());
106        }
107        assert_eq!(
108            states.len(),
109            1,
110            "count_distinct states must be single array"
111        );
112
113        let arr = as_list_array(&states[0])?;
114        arr.iter().try_for_each(|maybe_list| {
115            if let Some(list) = maybe_list {
116                let list = as_primitive_array::<T>(&list)?;
117                self.values.extend(list.values())
118            };
119            Ok(())
120        })
121    }
122
123    fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
124        Ok(ScalarValue::Int64(Some(self.values.len() as i64)))
125    }
126
127    fn size(&self) -> usize {
128        let num_elements = self.values.len();
129        let fixed_size = size_of_val(self) + size_of_val(&self.values);
130
131        estimate_memory_size::<T::Native>(num_elements, fixed_size).unwrap()
132    }
133}
134
135#[derive(Debug)]
136pub struct FloatDistinctCountAccumulator<T: ArrowPrimitiveType> {
137    values: GenericDistinctBuffer<T>,
138}
139
140impl<T: ArrowPrimitiveType> FloatDistinctCountAccumulator<T> {
141    pub fn new() -> Self {
142        Self {
143            values: GenericDistinctBuffer::new(T::DATA_TYPE),
144        }
145    }
146}
147
148impl<T: ArrowPrimitiveType> Default for FloatDistinctCountAccumulator<T> {
149    fn default() -> Self {
150        Self::new()
151    }
152}
153
154impl<T: ArrowPrimitiveType + Debug> Accumulator for FloatDistinctCountAccumulator<T> {
155    fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
156        self.values.state()
157    }
158
159    #[inline(never)]
160    fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
161        self.values.update_batch(values)
162    }
163
164    fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> {
165        self.values.merge_batch(states)
166    }
167
168    fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
169        Ok(ScalarValue::Int64(Some(self.values.values.len() as i64)))
170    }
171
172    fn size(&self) -> usize {
173        size_of_val(self) + self.values.size()
174    }
175}
176
177/// Optimized COUNT DISTINCT accumulator for u8 using a bool array.
178/// Uses 256 bytes to track all possible u8 values.
179#[derive(Debug)]
180pub struct BoolArray256DistinctCountAccumulator {
181    seen: [bool; 256],
182}
183
184impl BoolArray256DistinctCountAccumulator {
185    pub fn new() -> Self {
186        Self { seen: [false; 256] }
187    }
188
189    #[inline]
190    fn count(&self) -> i64 {
191        self.seen.iter().filter(|&&b| b).count() as i64
192    }
193}
194
195impl Default for BoolArray256DistinctCountAccumulator {
196    fn default() -> Self {
197        Self::new()
198    }
199}
200
201impl Accumulator for BoolArray256DistinctCountAccumulator {
202    #[inline(never)]
203    fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
204        if values.is_empty() {
205            return Ok(());
206        }
207
208        let arr = as_primitive_array::<arrow::datatypes::UInt8Type>(&values[0])?;
209        for value in arr.iter().flatten() {
210            self.seen[value as usize] = true;
211        }
212        Ok(())
213    }
214
215    fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> {
216        if states.is_empty() {
217            return Ok(());
218        }
219
220        let arr = as_list_array(&states[0])?;
221        arr.iter().try_for_each(|maybe_list| {
222            if let Some(list) = maybe_list {
223                let list = as_primitive_array::<arrow::datatypes::UInt8Type>(&list)?;
224                for value in list.values().iter() {
225                    self.seen[*value as usize] = true;
226                }
227            };
228            Ok(())
229        })
230    }
231
232    fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
233        let values: Vec<u8> = self
234            .seen
235            .iter()
236            .enumerate()
237            .filter_map(|(idx, &seen)| if seen { Some(idx as u8) } else { None })
238            .collect();
239
240        let arr = Arc::new(
241            PrimitiveArray::<arrow::datatypes::UInt8Type>::from_iter_values(values),
242        );
243        Ok(vec![
244            SingleRowListArrayBuilder::new(arr).build_list_scalar(),
245        ])
246    }
247
248    fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
249        Ok(ScalarValue::Int64(Some(self.count())))
250    }
251
252    fn size(&self) -> usize {
253        size_of_val(self) + 256
254    }
255}
256
257/// Optimized COUNT DISTINCT accumulator for i8 using a bool array.
258/// Uses 256 bytes to track all possible i8 values (mapped to 0..255).
259#[derive(Debug)]
260pub struct BoolArray256DistinctCountAccumulatorI8 {
261    seen: [bool; 256],
262}
263
264impl BoolArray256DistinctCountAccumulatorI8 {
265    pub fn new() -> Self {
266        Self { seen: [false; 256] }
267    }
268
269    #[inline]
270    fn count(&self) -> i64 {
271        self.seen.iter().filter(|&&b| b).count() as i64
272    }
273}
274
275impl Default for BoolArray256DistinctCountAccumulatorI8 {
276    fn default() -> Self {
277        Self::new()
278    }
279}
280
281impl Accumulator for BoolArray256DistinctCountAccumulatorI8 {
282    #[inline(never)]
283    fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
284        if values.is_empty() {
285            return Ok(());
286        }
287
288        let arr = as_primitive_array::<arrow::datatypes::Int8Type>(&values[0])?;
289        for value in arr.iter().flatten() {
290            self.seen[value as u8 as usize] = true;
291        }
292        Ok(())
293    }
294
295    fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> {
296        if states.is_empty() {
297            return Ok(());
298        }
299
300        let arr = as_list_array(&states[0])?;
301        arr.iter().try_for_each(|maybe_list| {
302            if let Some(list) = maybe_list {
303                let list = as_primitive_array::<arrow::datatypes::Int8Type>(&list)?;
304                for value in list.values().iter() {
305                    self.seen[*value as u8 as usize] = true;
306                }
307            };
308            Ok(())
309        })
310    }
311
312    fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
313        let values: Vec<i8> = self
314            .seen
315            .iter()
316            .enumerate()
317            .filter_map(
318                |(idx, &seen)| {
319                    if seen { Some(idx as u8 as i8) } else { None }
320                },
321            )
322            .collect();
323
324        let arr = Arc::new(
325            PrimitiveArray::<arrow::datatypes::Int8Type>::from_iter_values(values),
326        );
327        Ok(vec![
328            SingleRowListArrayBuilder::new(arr).build_list_scalar(),
329        ])
330    }
331
332    fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
333        Ok(ScalarValue::Int64(Some(self.count())))
334    }
335
336    fn size(&self) -> usize {
337        size_of_val(self) + 256
338    }
339}
340
341/// Optimized COUNT DISTINCT accumulator for u16 using a 65536-bit bitmap.
342/// Uses 8KB (1024 x u64) to track all possible u16 values.
343#[derive(Debug)]
344pub struct Bitmap65536DistinctCountAccumulator {
345    bitmap: Box<[u64; 1024]>,
346}
347
348impl Bitmap65536DistinctCountAccumulator {
349    pub fn new() -> Self {
350        Self {
351            bitmap: Box::new([0; 1024]),
352        }
353    }
354
355    #[inline]
356    fn set_bit(&mut self, value: u16) {
357        let word = (value / 64) as usize;
358        let bit = value % 64;
359        self.bitmap[word] |= 1u64 << bit;
360    }
361
362    #[inline]
363    fn count(&self) -> i64 {
364        self.bitmap.iter().map(|w| w.count_ones() as i64).sum()
365    }
366}
367
368impl Default for Bitmap65536DistinctCountAccumulator {
369    fn default() -> Self {
370        Self::new()
371    }
372}
373
374impl Accumulator for Bitmap65536DistinctCountAccumulator {
375    #[inline(never)]
376    fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
377        if values.is_empty() {
378            return Ok(());
379        }
380
381        let arr = as_primitive_array::<arrow::datatypes::UInt16Type>(&values[0])?;
382        for value in arr.iter().flatten() {
383            self.set_bit(value);
384        }
385        Ok(())
386    }
387
388    fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> {
389        if states.is_empty() {
390            return Ok(());
391        }
392
393        let arr = as_list_array(&states[0])?;
394        arr.iter().try_for_each(|maybe_list| {
395            if let Some(list) = maybe_list {
396                let list = as_primitive_array::<arrow::datatypes::UInt16Type>(&list)?;
397                for value in list.values().iter() {
398                    self.set_bit(*value);
399                }
400            };
401            Ok(())
402        })
403    }
404
405    fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
406        let mut values = Vec::new();
407        for (word_idx, &word) in self.bitmap.iter().enumerate() {
408            if word != 0 {
409                for bit in 0..64 {
410                    if (word & (1u64 << bit)) != 0 {
411                        values.push((word_idx as u16) * 64 + bit);
412                    }
413                }
414            }
415        }
416
417        let arr = Arc::new(
418            PrimitiveArray::<arrow::datatypes::UInt16Type>::from_iter_values(values),
419        );
420        Ok(vec![
421            SingleRowListArrayBuilder::new(arr).build_list_scalar(),
422        ])
423    }
424
425    fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
426        Ok(ScalarValue::Int64(Some(self.count())))
427    }
428
429    fn size(&self) -> usize {
430        size_of_val(self) + 8192
431    }
432}
433
434/// Optimized COUNT DISTINCT accumulator for i16 using a 65536-bit bitmap.
435/// Uses 8KB (1024 x u64) to track all possible i16 values (mapped to 0..65535).
436#[derive(Debug)]
437pub struct Bitmap65536DistinctCountAccumulatorI16 {
438    bitmap: Box<[u64; 1024]>,
439}
440
441impl Bitmap65536DistinctCountAccumulatorI16 {
442    pub fn new() -> Self {
443        Self {
444            bitmap: Box::new([0; 1024]),
445        }
446    }
447
448    #[inline]
449    fn set_bit(&mut self, value: i16) {
450        let idx = value as u16;
451        let word = (idx / 64) as usize;
452        let bit = idx % 64;
453        self.bitmap[word] |= 1u64 << bit;
454    }
455
456    #[inline]
457    fn count(&self) -> i64 {
458        self.bitmap.iter().map(|w| w.count_ones() as i64).sum()
459    }
460}
461
462impl Default for Bitmap65536DistinctCountAccumulatorI16 {
463    fn default() -> Self {
464        Self::new()
465    }
466}
467
468impl Accumulator for Bitmap65536DistinctCountAccumulatorI16 {
469    #[inline(never)]
470    fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
471        if values.is_empty() {
472            return Ok(());
473        }
474
475        let arr = as_primitive_array::<arrow::datatypes::Int16Type>(&values[0])?;
476        for value in arr.iter().flatten() {
477            self.set_bit(value);
478        }
479        Ok(())
480    }
481
482    fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> {
483        if states.is_empty() {
484            return Ok(());
485        }
486
487        let arr = as_list_array(&states[0])?;
488        arr.iter().try_for_each(|maybe_list| {
489            if let Some(list) = maybe_list {
490                let list = as_primitive_array::<arrow::datatypes::Int16Type>(&list)?;
491                for value in list.values().iter() {
492                    self.set_bit(*value);
493                }
494            };
495            Ok(())
496        })
497    }
498
499    fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
500        let mut values = Vec::new();
501        for (word_idx, &word) in self.bitmap.iter().enumerate() {
502            if word != 0 {
503                for bit in 0..64 {
504                    if (word & (1u64 << bit)) != 0 {
505                        values.push(((word_idx as u16) * 64 + bit) as i16);
506                    }
507                }
508            }
509        }
510
511        let arr = Arc::new(
512            PrimitiveArray::<arrow::datatypes::Int16Type>::from_iter_values(values),
513        );
514        Ok(vec![
515            SingleRowListArrayBuilder::new(arr).build_list_scalar(),
516        ])
517    }
518
519    fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
520        Ok(ScalarValue::Int64(Some(self.count())))
521    }
522
523    fn size(&self) -> usize {
524        size_of_val(self) + 8192
525    }
526}
527
528/// Optimized COUNT DISTINCT accumulator for `Boolean` using two flags.
529///
530/// Tracks whether `false` and `true` have been observed; nulls are skipped.
531/// Result is always 0, 1, or 2.
532#[derive(Debug)]
533pub struct BooleanDistinctCountAccumulator {
534    has_seen_false: bool,
535    has_seen_true: bool,
536}
537
538impl BooleanDistinctCountAccumulator {
539    pub fn new() -> Self {
540        Self {
541            has_seen_false: false,
542            has_seen_true: false,
543        }
544    }
545
546    #[inline]
547    fn seen_both(&self) -> bool {
548        self.has_seen_false && self.has_seen_true
549    }
550
551    #[inline]
552    fn count(&self) -> i64 {
553        (self.has_seen_false as u8 + self.has_seen_true as u8) as i64
554    }
555
556    /// Update flags from a `BooleanArray`, short-circuiting per-flag once set.
557    #[inline]
558    fn observe(&mut self, arr: &BooleanArray) {
559        if !self.has_seen_false && arr.has_false() {
560            self.has_seen_false = true;
561        }
562        if !self.has_seen_true && arr.has_true() {
563            self.has_seen_true = true;
564        }
565    }
566}
567
568impl Default for BooleanDistinctCountAccumulator {
569    fn default() -> Self {
570        Self::new()
571    }
572}
573
574impl Accumulator for BooleanDistinctCountAccumulator {
575    fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> {
576        if values.is_empty() || self.seen_both() {
577            return Ok(());
578        }
579
580        let arr = as_boolean_array(&values[0])?;
581        self.observe(arr);
582        Ok(())
583    }
584
585    fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> {
586        if states.is_empty() || self.seen_both() {
587            return Ok(());
588        }
589
590        let arr = as_list_array(&states[0])?;
591        arr.iter().try_for_each(|maybe_list| {
592            if self.seen_both() {
593                return Ok(());
594            }
595            if let Some(list) = maybe_list {
596                self.observe(as_boolean_array(&list)?);
597            };
598            Ok(())
599        })
600    }
601
602    fn state(&mut self) -> datafusion_common::Result<Vec<ScalarValue>> {
603        let mut values: Vec<bool> = Vec::with_capacity(2);
604        if self.has_seen_false {
605            values.push(false);
606        }
607        if self.has_seen_true {
608            values.push(true);
609        }
610
611        let arr = Arc::new(BooleanArray::from(values));
612        Ok(vec![
613            SingleRowListArrayBuilder::new(arr).build_list_scalar(),
614        ])
615    }
616
617    fn evaluate(&mut self) -> datafusion_common::Result<ScalarValue> {
618        Ok(ScalarValue::Int64(Some(self.count())))
619    }
620
621    fn size(&self) -> usize {
622        size_of_val(self)
623    }
624}
625
626#[cfg(test)]
627mod tests {
628    use super::*;
629    use arrow::array::Int64Array;
630    use arrow::datatypes::Int64Type;
631
632    #[test]
633    fn update_batch_null_free_fast_path_agrees_with_general_path() {
634        // The null-free fast path must produce the same distinct set as the
635        // general (validity-checking) path.
636        let dense: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3, 2, 1]));
637        let sparse: ArrayRef = Arc::new(Int64Array::from(vec![
638            Some(1),
639            None,
640            Some(2),
641            None,
642            Some(3),
643            Some(2),
644            Some(1),
645        ]));
646
647        let mut dense_acc =
648            PrimitiveDistinctCountAccumulator::<Int64Type>::new(&DataType::Int64);
649        dense_acc
650            .update_batch(std::slice::from_ref(&dense))
651            .unwrap();
652
653        let mut sparse_acc =
654            PrimitiveDistinctCountAccumulator::<Int64Type>::new(&DataType::Int64);
655        sparse_acc
656            .update_batch(std::slice::from_ref(&sparse))
657            .unwrap();
658
659        // Both should count the 3 distinct non-null values {1, 2, 3}.
660        assert_eq!(dense_acc.evaluate().unwrap(), ScalarValue::Int64(Some(3)));
661        assert_eq!(sparse_acc.evaluate().unwrap(), ScalarValue::Int64(Some(3)));
662    }
663}
datafusion_functions_aggregate_common/aggregate/count_distinct/native.rs

datafusion_functions_aggregate_common/aggregate/count_distinct/
native.rs
