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    // Initialize loop state.
367    let head_idx = validity
368        .first()
369        .vortex_expect("All null masks have been handled before");
370    let buffer = array.to_buffer::<T>();
371    let head = buffer[head_idx];
372
373    let mut loop_state = LoopState {
374        distinct_values: if count_distinct_values {
375            HashMap::with_capacity_and_hasher(array.len() / 2, FxBuildHasher)
376        } else {
377            HashMap::with_hasher(FxBuildHasher)
378        },
379        prev: head,
380        runs: 1,
381    };
382
383    let sliced = buffer.slice(head_idx..array.len());
384    let (chunks, remainder) = sliced.as_slice().as_chunks::<64>();
385    match validity.bit_buffer() {
386        AllOr::All => {
387            for chunk in chunks {
388                inner_loop_nonnull(chunk, count_distinct_values, &mut loop_state)
389            }
390            inner_loop_naive(
391                remainder,
392                count_distinct_values,
393                &BitBuffer::new_set(remainder.len()),
394                &mut loop_state,
395            );
396        }
397        AllOr::None => unreachable!("All invalid arrays have been handled before"),
398        AllOr::Some(v) => {
399            let mask = v.slice(head_idx..array.len());
400            let mut offset = 0;
401            for chunk in chunks {
402                let validity = mask.slice(offset..(offset + 64));
403                offset += 64;
404
405                match validity.true_count() {
406                    // All nulls -> no stats to update.
407                    0 => continue,
408                    // Inner loop for when validity check can be elided.
409                    64 => inner_loop_nonnull(chunk, count_distinct_values, &mut loop_state),
410                    // Inner loop for when we need to check validity.
411                    _ => inner_loop_nullable(
412                        chunk,
413                        count_distinct_values,
414                        &validity,
415                        &mut loop_state,
416                    ),
417                }
418            }
419            // Final iteration, run naive loop.
420            inner_loop_naive(
421                remainder,
422                count_distinct_values,
423                &mask.slice(offset..(offset + remainder.len())),
424                &mut loop_state,
425            );
426        }
427    }
428
429    let runs = loop_state.runs;
430
431    let array_ref = array.as_ref();
432    let min = array_ref
433        .statistics()
434        .compute_as::<T>(Stat::Min, ctx)
435        .vortex_expect("min should be computed");
436
437    let max = array_ref
438        .statistics()
439        .compute_as::<T>(Stat::Max, ctx)
440        .vortex_expect("max should be computed");
441
442    let distinct = count_distinct_values.then(|| {
443        let (&top_value, &top_count) = loop_state
444            .distinct_values
445            .iter()
446            .max_by_key(|&(_, &count)| count)
447            .vortex_expect("we know this is non-empty");
448
449        DistinctInfo {
450            distinct_count: u32::try_from(loop_state.distinct_values.len())
451                .vortex_expect("there are more than `u32::MAX` distinct values"),
452            most_frequent_value: top_value.0,
453            top_frequency: top_count,
454            distinct_values: loop_state.distinct_values,
455        }
456    });
457
458    let typed = TypedStats { min, max, distinct };
459
460    let null_count = u32::try_from(null_count)?;
461    let value_count = u32::try_from(value_count)?;
462
463    Ok(IntegerStats {
464        null_count,
465        value_count,
466        average_run_length: value_count / runs,
467        erased: typed.into(),
468    })
469}
470
471/// Internal loop state for integer stats computation.
472struct LoopState<T> {
473    /// The previous value seen.
474    prev: T,
475    /// The run count.
476    runs: u32,
477    /// The distinct values map.
478    distinct_values: HashMap<NativeValue<T>, u32, FxBuildHasher>,
479}
480
481/// Inner loop for non-null chunks of 64 values.
482#[allow(clippy::inline_always)]
483#[inline(always)]
484fn inner_loop_nonnull<T: IntegerPType>(
485    values: &[T; 64],
486    count_distinct_values: bool,
487    state: &mut LoopState<T>,
488) where
489    NativeValue<T>: Eq + Hash,
490{
491    for &value in values {
492        if count_distinct_values {
493            *state.distinct_values.entry(NativeValue(value)).or_insert(0) += 1;
494        }
495
496        if value != state.prev {
497            state.prev = value;
498            state.runs += 1;
499        }
500    }
501}
502
503/// Inner loop for nullable chunks of 64 values.
504#[allow(clippy::inline_always)]
505#[inline(always)]
506fn inner_loop_nullable<T: IntegerPType>(
507    values: &[T; 64],
508    count_distinct_values: bool,
509    is_valid: &BitBuffer,
510    state: &mut LoopState<T>,
511) where
512    NativeValue<T>: Eq + Hash,
513{
514    for (idx, &value) in values.iter().enumerate() {
515        if is_valid.value(idx) {
516            if count_distinct_values {
517                *state.distinct_values.entry(NativeValue(value)).or_insert(0) += 1;
518            }
519
520            if value != state.prev {
521                state.prev = value;
522                state.runs += 1;
523            }
524        }
525    }
526}
527
528/// Fallback inner loop for remainder values.
529#[allow(clippy::inline_always)]
530#[inline(always)]
531fn inner_loop_naive<T: IntegerPType>(
532    values: &[T],
533    count_distinct_values: bool,
534    is_valid: &BitBuffer,
535    state: &mut LoopState<T>,
536) where
537    NativeValue<T>: Eq + Hash,
538{
539    for (idx, &value) in values.iter().enumerate() {
540        if is_valid.value(idx) {
541            if count_distinct_values {
542                *state.distinct_values.entry(NativeValue(value)).or_insert(0) += 1;
543            }
544
545            if value != state.prev {
546                state.prev = value;
547                state.runs += 1;
548            }
549        }
550    }
551}
552
553#[cfg(test)]
554mod tests {
555    use std::iter;
556
557    use vortex_array::VortexSessionExecute;
558    use vortex_array::array_session;
559    use vortex_array::arrays::PrimitiveArray;
560    use vortex_array::validity::Validity;
561    use vortex_buffer::BitBuffer;
562    use vortex_buffer::Buffer;
563    use vortex_buffer::buffer;
564    use vortex_error::VortexResult;
565
566    use super::IntegerStats;
567    use super::typed_int_stats;
568
569    #[test]
570    fn test_naive_count_distinct_values() -> VortexResult<()> {
571        let mut ctx = array_session().create_execution_ctx();
572        let array = PrimitiveArray::new(buffer![217u8, 0], Validity::NonNullable);
573        let stats = typed_int_stats::<u8>(&array, true, &mut ctx)?;
574        assert_eq!(stats.distinct_count().unwrap(), 2);
575        Ok(())
576    }
577
578    #[test]
579    fn test_naive_count_distinct_values_nullable() -> VortexResult<()> {
580        let mut ctx = array_session().create_execution_ctx();
581        let array = PrimitiveArray::new(
582            buffer![217u8, 0],
583            Validity::from(BitBuffer::from(vec![true, false])),
584        );
585        let stats = typed_int_stats::<u8>(&array, true, &mut ctx)?;
586        assert_eq!(stats.distinct_count().unwrap(), 1);
587        Ok(())
588    }
589
590    #[test]
591    fn test_count_distinct_values() -> VortexResult<()> {
592        let mut ctx = array_session().create_execution_ctx();
593        let array = PrimitiveArray::new((0..128u8).collect::<Buffer<u8>>(), Validity::NonNullable);
594        let stats = typed_int_stats::<u8>(&array, true, &mut ctx)?;
595        assert_eq!(stats.distinct_count().unwrap(), 128);
596        Ok(())
597    }
598
599    #[test]
600    fn test_count_distinct_values_nullable() -> VortexResult<()> {
601        let mut ctx = array_session().create_execution_ctx();
602        let array = PrimitiveArray::new(
603            (0..128u8).collect::<Buffer<u8>>(),
604            Validity::from(BitBuffer::from_iter(
605                iter::repeat_n(vec![true, false], 64).flatten(),
606            )),
607        );
608        let stats = typed_int_stats::<u8>(&array, true, &mut ctx)?;
609        assert_eq!(stats.distinct_count().unwrap(), 64);
610        Ok(())
611    }
612
613    #[test]
614    fn test_integer_stats_leading_nulls() {
615        let mut ctx = array_session().create_execution_ctx();
616        let ints = PrimitiveArray::new(buffer![0, 1, 2], Validity::from_iter([false, true, true]));
617
618        let stats = IntegerStats::generate_opts(
619            &ints,
620            crate::stats::GenerateStatsOptions {
621                count_distinct_values: true,
622            },
623            &mut ctx,
624        );
625
626        assert_eq!(stats.value_count, 2);
627        assert_eq!(stats.null_count, 1);
628        assert_eq!(stats.average_run_length, 1);
629        assert_eq!(stats.distinct_count().unwrap(), 2);
630    }
631}
vortex_compressor/stats/integer.rs

vortex_compressor/stats/
integer.rs
