vortex_compressor/stats/
integer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Integer compression statistics.
5
6use std::hash::Hash;
7
8use num_traits::PrimInt;
9use rustc_hash::FxBuildHasher;
10use vortex_array::ExecutionCtx;
11use vortex_array::arrays::PrimitiveArray;
12use vortex_array::arrays::primitive::NativeValue;
13use vortex_array::dtype::IntegerPType;
14use vortex_array::expr::stats::Stat;
15use vortex_array::match_each_integer_ptype;
16use vortex_array::scalar::PValue;
17use vortex_array::scalar::Scalar;
18use vortex_buffer::BitBuffer;
19use vortex_error::VortexError;
20use vortex_error::VortexExpect;
21use vortex_error::VortexResult;
22use vortex_mask::AllOr;
23use vortex_utils::aliases::hash_map::HashMap;
24
25use super::GenerateStatsOptions;
26
27/// Information about the distinct values in an integer array.
28#[derive(Debug, Clone)]
29pub struct DistinctInfo<T> {
30    /// The unique values and their occurrences.
31    distinct_values: HashMap<NativeValue<T>, u32, FxBuildHasher>,
32    /// The count of unique values. This _must_ be non-zero.
33    distinct_count: u32,
34    /// The most frequent value.
35    most_frequent_value: T,
36    /// The number of times the most frequent value occurs.
37    top_frequency: u32,
38}
39
40impl<T> DistinctInfo<T> {
41    /// Returns a reference to the distinct values map.
42    pub fn distinct_values(&self) -> &HashMap<NativeValue<T>, u32, FxBuildHasher> {
43        &self.distinct_values
44    }
45}
46
47/// Typed statistics for a specific integer type.
48#[derive(Debug, Clone)]
49pub struct TypedStats<T> {
50    /// The minimum value.
51    min: T,
52    /// The maximum value.
53    max: T,
54    /// Distinct value information, or `None` if not computed.
55    distinct: Option<DistinctInfo<T>>,
56}
57
58impl<T> TypedStats<T> {
59    /// Returns the distinct value information, if computed.
60    pub fn distinct(&self) -> Option<&DistinctInfo<T>> {
61        self.distinct.as_ref()
62    }
63}
64
65impl<T> TypedStats<T> {
66    /// Get the count of distinct values, if we have computed it already.
67    fn distinct_count(&self) -> Option<u32> {
68        Some(self.distinct.as_ref()?.distinct_count)
69    }
70
71    /// Get the most commonly occurring value and its count, if we have computed it already.
72    fn most_frequent_value_and_count(&self) -> Option<(&T, u32)> {
73        let distinct = self.distinct.as_ref()?;
74        Some((&distinct.most_frequent_value, distinct.top_frequency))
75    }
76}
77
78/// Type-erased container for one of the [`TypedStats`] variants.
79///
80/// Building the `TypedStats` is considerably faster and cheaper than building a type-erased
81/// set of stats. We then perform a variety of access methods on them.
82#[derive(Clone, Debug)]
83pub enum ErasedStats {
84    /// Stats for `u8` arrays.
85    U8(TypedStats<u8>),
86    /// Stats for `u16` arrays.
87    U16(TypedStats<u16>),
88    /// Stats for `u32` arrays.
89    U32(TypedStats<u32>),
90    /// Stats for `u64` arrays.
91    U64(TypedStats<u64>),
92    /// Stats for `i8` arrays.
93    I8(TypedStats<i8>),
94    /// Stats for `i16` arrays.
95    I16(TypedStats<i16>),
96    /// Stats for `i32` arrays.
97    I32(TypedStats<i32>),
98    /// Stats for `i64` arrays.
99    I64(TypedStats<i64>),
100}
101
102impl ErasedStats {
103    /// Returns `true` if the minimum value is zero.
104    pub fn min_is_zero(&self) -> bool {
105        match &self {
106            ErasedStats::U8(x) => x.min == 0,
107            ErasedStats::U16(x) => x.min == 0,
108            ErasedStats::U32(x) => x.min == 0,
109            ErasedStats::U64(x) => x.min == 0,
110            ErasedStats::I8(x) => x.min == 0,
111            ErasedStats::I16(x) => x.min == 0,
112            ErasedStats::I32(x) => x.min == 0,
113            ErasedStats::I64(x) => x.min == 0,
114        }
115    }
116
117    /// Returns `true` if the minimum value is negative.
118    pub fn min_is_negative(&self) -> bool {
119        match &self {
120            ErasedStats::U8(_)
121            | ErasedStats::U16(_)
122            | ErasedStats::U32(_)
123            | ErasedStats::U64(_) => false,
124            ErasedStats::I8(x) => x.min < 0,
125            ErasedStats::I16(x) => x.min < 0,
126            ErasedStats::I32(x) => x.min < 0,
127            ErasedStats::I64(x) => x.min < 0,
128        }
129    }
130
131    /// Difference between max and min.
132    pub fn max_minus_min(&self) -> u64 {
133        match &self {
134            ErasedStats::U8(x) => (x.max - x.min) as u64,
135            ErasedStats::U16(x) => (x.max - x.min) as u64,
136            ErasedStats::U32(x) => (x.max - x.min) as u64,
137            ErasedStats::U64(x) => x.max - x.min,
138            ErasedStats::I8(x) => (x.max as i16 - x.min as i16) as u64,
139            ErasedStats::I16(x) => (x.max as i32 - x.min as i32) as u64,
140            ErasedStats::I32(x) => (x.max as i64 - x.min as i64) as u64,
141            ErasedStats::I64(x) => u64::try_from(x.max as i128 - x.min as i128)
142                .vortex_expect("max minus min result bigger than u64"),
143        }
144    }
145
146    /// Returns the ilog2 of the max value when transmuted to unsigned, or `None` if zero.
147    ///
148    /// This matches how BitPacking computes bit width: it reinterprets signed values as
149    /// unsigned (preserving bit pattern) and uses `leading_zeros`. For non-negative signed
150    /// values, the transmuted value equals the original value.
151    ///
152    /// This is used to determine if FOR encoding would reduce bit width compared to
153    /// direct BitPacking. If `max_ilog2() == max_minus_min_ilog2()`, FOR doesn't help.
154    pub fn max_ilog2(&self) -> Option<u32> {
155        match &self {
156            ErasedStats::U8(x) => x.max.checked_ilog2(),
157            ErasedStats::U16(x) => x.max.checked_ilog2(),
158            ErasedStats::U32(x) => x.max.checked_ilog2(),
159            ErasedStats::U64(x) => x.max.checked_ilog2(),
160            // Transmute signed to unsigned (bit pattern preserved) to match BitPacking behavior.
161            ErasedStats::I8(x) => (x.max as u8).checked_ilog2(),
162            ErasedStats::I16(x) => (x.max as u16).checked_ilog2(),
163            ErasedStats::I32(x) => (x.max as u32).checked_ilog2(),
164            ErasedStats::I64(x) => (x.max as u64).checked_ilog2(),
165        }
166    }
167
168    /// Get the count of distinct values, if we have computed it already.
169    pub fn distinct_count(&self) -> Option<u32> {
170        match &self {
171            ErasedStats::U8(x) => x.distinct_count(),
172            ErasedStats::U16(x) => x.distinct_count(),
173            ErasedStats::U32(x) => x.distinct_count(),
174            ErasedStats::U64(x) => x.distinct_count(),
175            ErasedStats::I8(x) => x.distinct_count(),
176            ErasedStats::I16(x) => x.distinct_count(),
177            ErasedStats::I32(x) => x.distinct_count(),
178            ErasedStats::I64(x) => x.distinct_count(),
179        }
180    }
181
182    /// Get the most commonly occurring value and its count.
183    pub fn most_frequent_value_and_count(&self) -> Option<(PValue, u32)> {
184        match &self {
185            ErasedStats::U8(x) => {
186                let (top_value, count) = x.most_frequent_value_and_count()?;
187                Some(((*top_value).into(), count))
188            }
189            ErasedStats::U16(x) => {
190                let (top_value, count) = x.most_frequent_value_and_count()?;
191                Some(((*top_value).into(), count))
192            }
193            ErasedStats::U32(x) => {
194                let (top_value, count) = x.most_frequent_value_and_count()?;
195                Some(((*top_value).into(), count))
196            }
197            ErasedStats::U64(x) => {
198                let (top_value, count) = x.most_frequent_value_and_count()?;
199                Some(((*top_value).into(), count))
200            }
201            ErasedStats::I8(x) => {
202                let (top_value, count) = x.most_frequent_value_and_count()?;
203                Some(((*top_value).into(), count))
204            }
205            ErasedStats::I16(x) => {
206                let (top_value, count) = x.most_frequent_value_and_count()?;
207                Some(((*top_value).into(), count))
208            }
209            ErasedStats::I32(x) => {
210                let (top_value, count) = x.most_frequent_value_and_count()?;
211                Some(((*top_value).into(), count))
212            }
213            ErasedStats::I64(x) => {
214                let (top_value, count) = x.most_frequent_value_and_count()?;
215                Some(((*top_value).into(), count))
216            }
217        }
218    }
219}
220
221/// Implements `From<TypedStats<$T>>` for [`ErasedStats`].
222macro_rules! impl_from_typed {
223    ($T:ty, $variant:path) => {
224        impl From<TypedStats<$T>> for ErasedStats {
225            fn from(typed: TypedStats<$T>) -> Self {
226                $variant(typed)
227            }
228        }
229    };
230}
231
232impl_from_typed!(u8, ErasedStats::U8);
233impl_from_typed!(u16, ErasedStats::U16);
234impl_from_typed!(u32, ErasedStats::U32);
235impl_from_typed!(u64, ErasedStats::U64);
236impl_from_typed!(i8, ErasedStats::I8);
237impl_from_typed!(i16, ErasedStats::I16);
238impl_from_typed!(i32, ErasedStats::I32);
239impl_from_typed!(i64, ErasedStats::I64);
240
241/// Array of integers and relevant stats for compression.
242#[derive(Clone, Debug)]
243pub struct IntegerStats {
244    /// Cache for `validity.false_count()`.
245    null_count: u32,
246    /// Cache for `validity.true_count()`.
247    value_count: u32,
248    /// The average run length.
249    average_run_length: u32,
250    /// Type-erased typed statistics.
251    erased: ErasedStats,
252}
253
254impl IntegerStats {
255    /// Generates stats, returning an error on failure.
256    fn generate_opts_fallible(
257        input: &PrimitiveArray,
258        opts: GenerateStatsOptions,
259        ctx: &mut ExecutionCtx,
260    ) -> VortexResult<Self> {
261        match_each_integer_ptype!(input.ptype(), |T| {
262            typed_int_stats::<T>(input, opts.count_distinct_values, ctx)
263        })
264    }
265
266    /// Get the count of distinct values, if we have computed it already.
267    pub fn distinct_count(&self) -> Option<u32> {
268        self.erased.distinct_count()
269    }
270
271    /// Get the most commonly occurring value and its count, if we have computed it already.
272    pub fn most_frequent_value_and_count(&self) -> Option<(PValue, u32)> {
273        self.erased.most_frequent_value_and_count()
274    }
275}
276
277impl IntegerStats {
278    /// Generates stats with default options.
279    pub fn generate(input: &PrimitiveArray, ctx: &mut ExecutionCtx) -> Self {
280        Self::generate_opts(input, GenerateStatsOptions::default(), ctx)
281    }
282
283    /// Generates stats with provided options.
284    pub fn generate_opts(
285        input: &PrimitiveArray,
286        opts: GenerateStatsOptions,
287        ctx: &mut ExecutionCtx,
288    ) -> Self {
289        Self::generate_opts_fallible(input, opts, ctx)
290            .vortex_expect("IntegerStats::generate_opts should not fail")
291    }
292
293    /// Returns the number of null values.
294    pub fn null_count(&self) -> u32 {
295        self.null_count
296    }
297
298    /// Returns the number of non-null values.
299    pub fn value_count(&self) -> u32 {
300        self.value_count
301    }
302
303    /// Returns the average run length.
304    pub fn average_run_length(&self) -> u32 {
305        self.average_run_length
306    }
307
308    /// Returns the type-erased typed statistics.
309    pub fn erased(&self) -> &ErasedStats {
310        &self.erased
311    }
312}
313
314/// Computes typed integer statistics for a specific integer type.
315fn typed_int_stats<T>(
316    array: &PrimitiveArray,
317    count_distinct_values: bool,
318    ctx: &mut ExecutionCtx,
319) -> VortexResult<IntegerStats>
320where
321    T: IntegerPType + PrimInt + for<'a> TryFrom<&'a Scalar, Error = VortexError>,
322    TypedStats<T>: Into<ErasedStats>,
323    NativeValue<T>: Eq + Hash,
324{
325    // Special case: empty array.
326    if array.is_empty() {
327        return Ok(IntegerStats {
328            null_count: 0,
329            value_count: 0,
330            average_run_length: 0,
331            erased: TypedStats {
332                min: T::max_value(),
333                max: T::min_value(),
334                distinct: None,
335            }
336            .into(),
337        });
338    }
339
340    if array.all_invalid(ctx)? {
341        return Ok(IntegerStats {
342            null_count: u32::try_from(array.len())?,
343            value_count: 0,
344            average_run_length: 0,
345            erased: TypedStats {
346                min: T::max_value(),
347                max: T::min_value(),
348                distinct: Some(DistinctInfo {
349                    distinct_values: HashMap::with_capacity_and_hasher(0, FxBuildHasher),
350                    distinct_count: 0,
351                    most_frequent_value: T::zero(),
352                    top_frequency: 0,
353                }),
354            }
355            .into(),
356        });
357    }
358
359    let validity = array
360        .as_ref()
361        .validity()?
362        .execute_mask(array.as_ref().len(), ctx)?;
363    let null_count = validity.false_count();
364    let value_count = validity.true_count();
365
366    let array_ref = array.as_ref();
367    let min = array_ref
368        .statistics()
369        .compute_as::<T>(Stat::Min, ctx)
370        .vortex_expect("min should be computed");
371
372    let max = array_ref
373        .statistics()
374        .compute_as::<T>(Stat::Max, ctx)
375        .vortex_expect("max should be computed");
376
377    // Initialize loop state.
378    let head_idx = validity
379        .first()
380        .vortex_expect("All null masks have been handled before");
381    let buffer = array.to_buffer::<T>();
382    let head = buffer[head_idx];
383
384    let mut loop_state = LoopState {
385        distinct: if count_distinct_values {
386            DistinctCounter::new(min, max, array.len())
387        } else {
388            DistinctCounter::Off
389        },
390        prev: head,
391        pending: 0,
392        runs: 1,
393    };
394
395    let sliced = buffer.slice(head_idx..array.len());
396    let (chunks, remainder) = sliced.as_slice().as_chunks::<64>();
397    match validity.bit_buffer() {
398        AllOr::All => {
399            for chunk in chunks {
400                inner_loop_nonnull(chunk, count_distinct_values, &mut loop_state)
401            }
402            inner_loop_naive(
403                remainder,
404                count_distinct_values,
405                &BitBuffer::new_set(remainder.len()),
406                &mut loop_state,
407            );
408        }
409        AllOr::None => unreachable!("All invalid arrays have been handled before"),
410        AllOr::Some(v) => {
411            let mask = v.slice(head_idx..array.len());
412            let mut offset = 0;
413            for chunk in chunks {
414                let validity = mask.slice(offset..(offset + 64));
415                offset += 64;
416
417                match validity.true_count() {
418                    // All nulls -> no stats to update.
419                    0 => continue,
420                    // Inner loop for when validity check can be elided.
421                    64 => inner_loop_nonnull(chunk, count_distinct_values, &mut loop_state),
422                    // Inner loop for when we need to check validity.
423                    _ => inner_loop_nullable(
424                        chunk,
425                        count_distinct_values,
426                        &validity,
427                        &mut loop_state,
428                    ),
429                }
430            }
431            // Final iteration, run naive loop.
432            inner_loop_naive(
433                remainder,
434                count_distinct_values,
435                &mask.slice(offset..(offset + remainder.len())),
436                &mut loop_state,
437            );
438        }
439    }
440
441    if count_distinct_values {
442        loop_state.flush();
443    }
444
445    let runs = loop_state.runs;
446    let distinct = loop_state.distinct.finish();
447
448    let typed = TypedStats { min, max, distinct };
449
450    let null_count = u32::try_from(null_count)?;
451    let value_count = u32::try_from(value_count)?;
452
453    Ok(IntegerStats {
454        null_count,
455        value_count,
456        average_run_length: value_count / runs,
457        erased: typed.into(),
458    })
459}
460
461/// Value ranges up to this many values may be counted in a dense array instead of a hash map.
462const DENSE_DISTINCT_MAX_RANGE: usize = 1 << 16;
463
464/// Value ranges up to this many values are always counted in a dense array, regardless of the array
465/// length. This covers every `u8` and `i8` array.
466const DENSE_DISTINCT_ALWAYS_RANGE: usize = 1 << 8;
467
468/// Accumulates the occurrences of each distinct valid value.
469enum DistinctCounter<T> {
470    /// Distinct values are not being counted.
471    Off,
472    /// Counts indexed by `value - min`, used when the value range is narrow. This avoids hashing
473    /// entirely.
474    Dense {
475        /// The minimum valid value.
476        min: T,
477        /// `min` cast to `usize`. Casting preserves values modulo `2^usize::BITS`, so the wrapping
478        /// difference of two cast values is the exact difference of any two values in the range.
479        min_index: usize,
480        /// The number of occurrences of `min + index`.
481        counts: Vec<u32>,
482    },
483    /// Counts keyed by value.
484    Hashed(HashMap<NativeValue<T>, u32, FxBuildHasher>),
485}
486
487impl<T: IntegerPType> DistinctCounter<T>
488where
489    NativeValue<T>: Eq + Hash,
490{
491    /// Chooses a dense counter when the valid values span a narrow range, otherwise a hash map.
492    fn new(min: T, max: T, len: usize) -> Self {
493        let range_len = match (min.to_i128(), max.to_i128()) {
494            (Some(min), Some(max)) => usize::try_from(max - min + 1).unwrap_or(usize::MAX),
495            _ => usize::MAX,
496        };
497
498        // A dense counter is only worthwhile when it is not much larger than the array itself.
499        if range_len <= DENSE_DISTINCT_ALWAYS_RANGE
500            || (range_len <= DENSE_DISTINCT_MAX_RANGE && range_len <= len)
501        {
502            Self::Dense {
503                min,
504                min_index: min.as_(),
505                counts: vec![0; range_len],
506            }
507        } else {
508            Self::Hashed(HashMap::with_capacity_and_hasher(len / 2, FxBuildHasher))
509        }
510    }
511
512    /// Records `count` occurrences of a valid value.
513    #[allow(clippy::inline_always)]
514    #[inline(always)]
515    fn add(&mut self, value: T, count: u32) {
516        match self {
517            Self::Off => {}
518            Self::Dense {
519                min_index, counts, ..
520            } => counts[value.as_().wrapping_sub(*min_index)] += count,
521            Self::Hashed(distinct_values) => {
522                *distinct_values.entry(NativeValue(value)).or_insert(0) += count
523            }
524        }
525    }
526
527    /// Returns the distinct value information, or `None` if counting was off.
528    fn finish(self) -> Option<DistinctInfo<T>> {
529        let distinct_values: HashMap<NativeValue<T>, u32, FxBuildHasher> = match self {
530            Self::Off => return None,
531            Self::Dense { min, counts, .. } => {
532                let min = min.to_i128().vortex_expect("integers fit in i128");
533                counts
534                    .iter()
535                    .enumerate()
536                    .filter(|&(_, &count)| count > 0)
537                    .map(|(index, &count)| {
538                        let value = <T as num_traits::NumCast>::from(min + index as i128)
539                            .vortex_expect("values between min and max fit in the type");
540                        (NativeValue(value), count)
541                    })
542                    .collect()
543            }
544            Self::Hashed(distinct_values) => distinct_values,
545        };
546
547        let (&top_value, &top_count) = distinct_values
548            .iter()
549            .max_by_key(|&(_, &count)| count)
550            .vortex_expect("we know this is non-empty");
551
552        Some(DistinctInfo {
553            distinct_count: u32::try_from(distinct_values.len())
554                .vortex_expect("there are more than `u32::MAX` distinct values"),
555            most_frequent_value: top_value.0,
556            top_frequency: top_count,
557            distinct_values,
558        })
559    }
560}
561
562/// Internal loop state for integer stats computation.
563struct LoopState<T> {
564    /// The previous value seen.
565    prev: T,
566    /// Occurrences of `prev` in the current run not yet added to `distinct`.
567    pending: u32,
568    /// The run count.
569    runs: u32,
570    /// The distinct values counter.
571    distinct: DistinctCounter<T>,
572}
573
574impl<T: IntegerPType> LoopState<T>
575where
576    NativeValue<T>: Eq + Hash,
577{
578    /// Adds the pending occurrences of `prev` to the distinct values counter.
579    #[allow(clippy::inline_always)]
580    #[inline(always)]
581    fn flush(&mut self) {
582        self.distinct.add(self.prev, self.pending);
583        self.pending = 0;
584    }
585
586    /// Records one value, hashing only when the value differs from `prev`.
587    #[allow(clippy::inline_always)]
588    #[inline(always)]
589    fn push<const COUNT_DISTINCT_VALUES: bool>(&mut self, value: T) {
590        if value != self.prev {
591            if COUNT_DISTINCT_VALUES {
592                self.flush();
593            }
594            self.prev = value;
595            self.runs += 1;
596        }
597        if COUNT_DISTINCT_VALUES {
598            self.pending += 1;
599        }
600    }
601}
602
603/// Inner loop for non-null chunks of 64 values.
604#[allow(clippy::inline_always)]
605#[inline(always)]
606fn inner_loop_nonnull<T: IntegerPType>(
607    values: &[T; 64],
608    count_distinct_values: bool,
609    state: &mut LoopState<T>,
610) where
611    NativeValue<T>: Eq + Hash,
612{
613    if count_distinct_values {
614        inner_loop_nonnull_impl::<T, true>(values, state);
615    } else {
616        inner_loop_nonnull_impl::<T, false>(values, state);
617    }
618}
619
620/// Processes one non-null chunk, monomorphized on whether distinct values are counted.
621#[allow(clippy::inline_always)]
622#[inline(always)]
623fn inner_loop_nonnull_impl<T: IntegerPType, const COUNT_DISTINCT_VALUES: bool>(
624    values: &[T; 64],
625    state: &mut LoopState<T>,
626) where
627    NativeValue<T>: Eq + Hash,
628{
629    // Branch-free count of value changes, including the change from the previous chunk. At
630    // most 64, so a `u8` accumulator lets the comparison use full-width byte lanes.
631    let transitions = u8::from(values[0] != state.prev)
632        + values
633            .iter()
634            .zip(&values[1..])
635            .map(|(a, b)| u8::from(a != b))
636            .sum::<u8>();
637
638    if COUNT_DISTINCT_VALUES {
639        if transitions == 0 {
640            state.pending += 64;
641            return;
642        }
643        for &value in values {
644            if value != state.prev {
645                state.flush();
646                state.prev = value;
647            }
648            state.pending += 1;
649        }
650    }
651
652    state.runs += u32::from(transitions);
653    state.prev = values[63];
654}
655
656/// Inner loop for nullable chunks of 64 values.
657#[allow(clippy::inline_always)]
658#[inline(always)]
659fn inner_loop_nullable<T: IntegerPType>(
660    values: &[T; 64],
661    count_distinct_values: bool,
662    is_valid: &BitBuffer,
663    state: &mut LoopState<T>,
664) where
665    NativeValue<T>: Eq + Hash,
666{
667    if count_distinct_values {
668        inner_loop_masked::<T, true>(values, is_valid, state);
669    } else {
670        inner_loop_masked::<T, false>(values, is_valid, state);
671    }
672}
673
674/// Fallback inner loop for remainder values.
675#[allow(clippy::inline_always)]
676#[inline(always)]
677fn inner_loop_naive<T: IntegerPType>(
678    values: &[T],
679    count_distinct_values: bool,
680    is_valid: &BitBuffer,
681    state: &mut LoopState<T>,
682) where
683    NativeValue<T>: Eq + Hash,
684{
685    if count_distinct_values {
686        inner_loop_masked::<T, true>(values, is_valid, state);
687    } else {
688        inner_loop_masked::<T, false>(values, is_valid, state);
689    }
690}
691
692/// Processes values with a validity mask, skipping nulls without breaking runs.
693#[allow(clippy::inline_always)]
694#[inline(always)]
695fn inner_loop_masked<T: IntegerPType, const COUNT_DISTINCT_VALUES: bool>(
696    values: &[T],
697    is_valid: &BitBuffer,
698    state: &mut LoopState<T>,
699) where
700    NativeValue<T>: Eq + Hash,
701{
702    for (idx, &value) in values.iter().enumerate() {
703        if is_valid.value(idx) {
704            state.push::<COUNT_DISTINCT_VALUES>(value);
705        }
706    }
707}
708
709#[cfg(test)]
710mod tests {
711    use std::iter;
712
713    use rstest::rstest;
714    use vortex_array::VortexSessionExecute;
715    use vortex_array::array_session;
716    use vortex_array::arrays::PrimitiveArray;
717    use vortex_array::validity::Validity;
718    use vortex_buffer::BitBuffer;
719    use vortex_buffer::Buffer;
720    use vortex_buffer::buffer;
721    use vortex_error::VortexResult;
722    use vortex_utils::aliases::hash_map::HashMap;
723
724    use super::ErasedStats;
725    use super::IntegerStats;
726    use super::typed_int_stats;
727
728    #[test]
729    fn test_naive_count_distinct_values() -> VortexResult<()> {
730        let mut ctx = array_session().create_execution_ctx();
731        let array = PrimitiveArray::new(buffer![217u8, 0], Validity::NonNullable);
732        let stats = typed_int_stats::<u8>(&array, true, &mut ctx)?;
733        assert_eq!(stats.distinct_count().unwrap(), 2);
734        Ok(())
735    }
736
737    #[test]
738    fn test_naive_count_distinct_values_nullable() -> VortexResult<()> {
739        let mut ctx = array_session().create_execution_ctx();
740        let array = PrimitiveArray::new(
741            buffer![217u8, 0],
742            Validity::from(BitBuffer::from(vec![true, false])),
743        );
744        let stats = typed_int_stats::<u8>(&array, true, &mut ctx)?;
745        assert_eq!(stats.distinct_count().unwrap(), 1);
746        Ok(())
747    }
748
749    #[test]
750    fn test_count_distinct_values() -> VortexResult<()> {
751        let mut ctx = array_session().create_execution_ctx();
752        let array = PrimitiveArray::new((0..128u8).collect::<Buffer<u8>>(), Validity::NonNullable);
753        let stats = typed_int_stats::<u8>(&array, true, &mut ctx)?;
754        assert_eq!(stats.distinct_count().unwrap(), 128);
755        Ok(())
756    }
757
758    #[test]
759    fn test_count_distinct_values_nullable() -> VortexResult<()> {
760        let mut ctx = array_session().create_execution_ctx();
761        let array = PrimitiveArray::new(
762            (0..128u8).collect::<Buffer<u8>>(),
763            Validity::from(BitBuffer::from_iter(
764                iter::repeat_n(vec![true, false], 64).flatten(),
765            )),
766        );
767        let stats = typed_int_stats::<u8>(&array, true, &mut ctx)?;
768        assert_eq!(stats.distinct_count().unwrap(), 64);
769        Ok(())
770    }
771
772    #[test]
773    fn dense_distinct_signed_full_range() -> VortexResult<()> {
774        let mut ctx = array_session().create_execution_ctx();
775        let values: Buffer<i8> = (i8::MIN..=i8::MAX)
776            .chain(i8::MIN..=i8::MAX)
777            .chain(iter::once(i8::MAX))
778            .collect();
779        let array = PrimitiveArray::new(values, Validity::NonNullable);
780
781        let stats = typed_int_stats::<i8>(&array, true, &mut ctx)?;
782        assert_eq!(stats.distinct_count(), Some(256));
783        assert_eq!(
784            stats.most_frequent_value_and_count(),
785            Some((i8::MAX.into(), 3))
786        );
787        Ok(())
788    }
789
790    #[test]
791    fn dense_distinct_ignores_values_under_nulls() -> VortexResult<()> {
792        let mut ctx = array_session().create_execution_ctx();
793        let array = PrimitiveArray::new(
794            buffer![1000i32, 1001, 1000, 5_000_000, 1003],
795            Validity::from_iter([true, true, true, false, true]),
796        );
797
798        let stats = typed_int_stats::<i32>(&array, true, &mut ctx)?;
799        assert_eq!(stats.distinct_count(), Some(3));
800        assert_eq!(
801            stats.most_frequent_value_and_count(),
802            Some((1000i32.into(), 2))
803        );
804        Ok(())
805    }
806
807    #[test]
808    fn wide_range_distinct() -> VortexResult<()> {
809        let mut ctx = array_session().create_execution_ctx();
810        let array = PrimitiveArray::new(buffer![0u64, u64::MAX, 0], Validity::NonNullable);
811
812        let stats = typed_int_stats::<u64>(&array, true, &mut ctx)?;
813        assert_eq!(stats.distinct_count(), Some(2));
814        assert_eq!(
815            stats.most_frequent_value_and_count(),
816            Some((0u64.into(), 2))
817        );
818        Ok(())
819    }
820
821    #[test]
822    fn test_integer_stats_leading_nulls() {
823        let mut ctx = array_session().create_execution_ctx();
824        let ints = PrimitiveArray::new(buffer![0, 1, 2], Validity::from_iter([false, true, true]));
825
826        let stats = IntegerStats::generate_opts(
827            &ints,
828            crate::stats::GenerateStatsOptions {
829                count_distinct_values: true,
830            },
831            &mut ctx,
832        );
833
834        assert_eq!(stats.value_count, 2);
835        assert_eq!(stats.null_count, 1);
836        assert_eq!(stats.average_run_length, 1);
837        assert_eq!(stats.distinct_count().unwrap(), 2);
838    }
839
840    /// Returns the per-value counts and the run count of the valid values, computed naively.
841    fn naive_stats(values: &[u32], valid: &[bool]) -> (HashMap<u32, u32>, u32) {
842        let mut counts = HashMap::default();
843        let mut runs = 0;
844        let mut prev = None;
845        for (&value, _) in values.iter().zip(valid).filter(|(_, ok)| **ok) {
846            *counts.entry(value).or_insert(0) += 1;
847            if prev.replace(value) != Some(value) {
848                runs += 1;
849            }
850        }
851        (counts, runs)
852    }
853
854    /// Checks the chunked loops against a naive reference. The run lengths cover aligned
855    /// constant chunks (64), runs that cross chunk boundaries (3, 100), and chunks where every
856    /// value changes (1).
857    #[rstest]
858    fn test_matches_naive_reference(
859        #[values(1, 63, 64, 65, 64 * 20 + 17)] len: u32,
860        #[values(1, 3, 64, 100)] run_len: u32,
861        #[values(None, Some(97), Some(3))] null_every: Option<u32>,
862    ) -> VortexResult<()> {
863        let values: Vec<u32> = (0..len).map(|i| (i / run_len) % 16).collect();
864        let valid: Vec<bool> = (0..len)
865            .map(|i| null_every.is_none_or(|n| i % n != 0))
866            .collect();
867        let (expected, runs) = naive_stats(&values, &valid);
868
869        let validity = match null_every {
870            None => Validity::NonNullable,
871            Some(_) => Validity::from(BitBuffer::from(valid)),
872        };
873        let array = PrimitiveArray::new(Buffer::from(values), validity);
874        let mut ctx = array_session().create_execution_ctx();
875        let stats = typed_int_stats::<u32>(&array, true, &mut ctx)?;
876
877        let ErasedStats::U32(typed) = stats.erased() else {
878            unreachable!()
879        };
880        let actual: HashMap<u32, u32> = typed
881            .distinct()
882            .map(|d| d.distinct_values().iter().map(|(k, &c)| (k.0, c)).collect())
883            .unwrap_or_default();
884        assert_eq!(actual, expected);
885        let value_count: u32 = expected.values().sum();
886        assert_eq!(stats.value_count, value_count);
887        if value_count > 0 {
888            assert_eq!(stats.average_run_length, value_count / runs);
889        }
890        Ok(())
891    }
892}
vortex_compressor/stats/integer.rs

vortex_compressor/stats/
integer.rs
