vortex_array/stats/
stats_set.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::type_name;
5use std::fmt::Debug;
6
7use enum_iterator::all;
8use num_traits::CheckedAdd;
9use smallvec::SmallVec;
10use smallvec::smallvec;
11use vortex_error::VortexError;
12use vortex_error::VortexExpect;
13use vortex_error::VortexResult;
14use vortex_error::vortex_err;
15use vortex_error::vortex_panic;
16
17use crate::dtype::DType;
18use crate::expr::stats::IsConstant;
19use crate::expr::stats::IsSorted;
20use crate::expr::stats::IsStrictSorted;
21use crate::expr::stats::Max;
22use crate::expr::stats::Min;
23use crate::expr::stats::NaNCount;
24use crate::expr::stats::NullCount;
25use crate::expr::stats::Precision;
26use crate::expr::stats::Stat;
27use crate::expr::stats::StatBound;
28use crate::expr::stats::StatType;
29use crate::expr::stats::StatsProvider;
30use crate::expr::stats::StatsProviderExt;
31use crate::expr::stats::Sum;
32use crate::expr::stats::UncompressedSizeInBytes;
33use crate::scalar::Scalar;
34use crate::scalar::ScalarValue;
35
36/// Type of the SmallVec stored inside StatsSet
37pub type StatsArray = [(Stat, Precision<ScalarValue>); 4];
38
39#[derive(Default, Debug, Clone)]
40pub struct StatsSet {
41    values: SmallVec<StatsArray>,
42}
43
44impl StatsSet {
45    /// Create new StatSet without validating uniqueness of all the entries
46    ///
47    /// # Safety
48    ///
49    /// This method will not panic or trigger UB, but may lead to duplicate stats being stored.
50    pub unsafe fn new_unchecked(values: SmallVec<StatsArray>) -> Self {
51        Self { values }
52    }
53
54    /// Create StatsSet from single stat and value
55    pub fn of(stat: Stat, value: Precision<ScalarValue>) -> Self {
56        Self {
57            values: smallvec![(stat, value)],
58        }
59    }
60
61    /// Wrap stats set with a dtype for mutable typed scalar access
62    pub fn as_mut_typed_ref<'a, 'b>(&'a mut self, dtype: &'b DType) -> MutTypedStatsSetRef<'a, 'b> {
63        MutTypedStatsSetRef {
64            values: self,
65            dtype,
66        }
67    }
68
69    /// Wrap stats set with a dtype for typed scalar access
70    pub fn as_typed_ref<'a, 'b>(&'a self, dtype: &'b DType) -> TypedStatsSetRef<'a, 'b> {
71        TypedStatsSetRef {
72            values: self,
73            dtype,
74        }
75    }
76}
77
78// Getters and setters for individual stats.
79impl StatsSet {
80    /// Set the stat `stat` to `value`.
81    pub fn set(&mut self, stat: Stat, value: Precision<ScalarValue>) {
82        if let Some(existing) = self.values.iter_mut().find(|(s, _)| *s == stat) {
83            *existing = (stat, value);
84        } else {
85            self.values.push((stat, value));
86        }
87    }
88
89    /// Clear the stat `stat` from the set.
90    pub fn clear(&mut self, stat: Stat) {
91        self.values.retain(|(s, _)| *s != stat);
92    }
93
94    /// Only keep given stats
95    pub fn retain_only(&mut self, stats: &[Stat]) {
96        self.values.retain(|(s, _)| stats.contains(s));
97    }
98
99    /// Iterate over the statistic names and values in-place.
100    ///
101    /// See [Iterator].
102    pub fn iter(&self) -> impl Iterator<Item = &(Stat, Precision<ScalarValue>)> {
103        self.values.iter()
104    }
105
106    /// Get value for a given stat
107    pub fn get(&self, stat: Stat) -> Precision<ScalarValue> {
108        self.get_ref(stat).map(ScalarValue::clone)
109    }
110
111    /// Get a reference to the value for a given stat.
112    fn get_ref(&self, stat: Stat) -> Precision<&ScalarValue> {
113        self.values
114            .iter()
115            .find(|(s, _)| *s == stat)
116            .map(|(_, v)| v.as_ref())
117            .unwrap_or(Precision::Absent)
118    }
119
120    /// Get the value for a count stat, such as [`Stat::NullCount`], without building a [`Scalar`].
121    pub(super) fn get_usize(&self, stat: Stat) -> Precision<usize> {
122        self.get_ref(stat).map(|v| {
123            usize::try_from(v).unwrap_or_else(|err| {
124                vortex_panic!(
125                    err,
126                    "Failed to get stat {} as {}",
127                    stat,
128                    type_name::<usize>()
129                )
130            })
131        })
132    }
133
134    /// Get the value for a boolean stat, such as [`Stat::IsSorted`], without building a
135    /// [`Scalar`].
136    pub(super) fn get_bool(&self, stat: Stat) -> Precision<bool> {
137        self.get_ref(stat).map(|v| match v {
138            ScalarValue::Bool(b) => *b,
139            _ => vortex_panic!("Failed to get stat {} as bool, found {}", stat, v),
140        })
141    }
142
143    /// Length of the stats set
144    pub fn len(&self) -> usize {
145        self.values.len()
146    }
147
148    /// Check whether the statset is empty
149    pub fn is_empty(&self) -> bool {
150        self.values.is_empty()
151    }
152
153    /// Get scalar value of a given dtype
154    pub fn get_as<T: for<'a> TryFrom<&'a Scalar, Error = VortexError>>(
155        &self,
156        stat: Stat,
157        dtype: &DType,
158    ) -> Precision<T> {
159        self.get(stat).map(|v| {
160            T::try_from(
161                &Scalar::try_new(dtype.clone(), Some(v))
162                    .vortex_expect("failed to construct a scalar statistic"),
163            )
164            .unwrap_or_else(|err| {
165                vortex_panic!(err, "Failed to get stat {} as {}", stat, type_name::<T>())
166            })
167        })
168    }
169}
170
171// StatSetIntoIter just exists to protect current implementation from exposure on the public API.
172
173/// Owned iterator over the stats.
174///
175/// See [IntoIterator].
176pub struct StatsSetIntoIter(smallvec::IntoIter<StatsArray>);
177
178impl Iterator for StatsSetIntoIter {
179    type Item = (Stat, Precision<ScalarValue>);
180
181    fn next(&mut self) -> Option<Self::Item> {
182        self.0.next()
183    }
184}
185
186impl IntoIterator for StatsSet {
187    type Item = (Stat, Precision<ScalarValue>);
188    type IntoIter = StatsSetIntoIter;
189
190    fn into_iter(self) -> Self::IntoIter {
191        StatsSetIntoIter(self.values.into_iter())
192    }
193}
194
195impl FromIterator<(Stat, Precision<ScalarValue>)> for StatsSet {
196    fn from_iter<T: IntoIterator<Item = (Stat, Precision<ScalarValue>)>>(iter: T) -> Self {
197        let iter = iter.into_iter();
198
199        let mut this = Self {
200            values: SmallVec::new(),
201        };
202        this.extend(iter);
203        this
204    }
205}
206
207impl Extend<(Stat, Precision<ScalarValue>)> for StatsSet {
208    #[inline]
209    fn extend<T: IntoIterator<Item = (Stat, Precision<ScalarValue>)>>(&mut self, iter: T) {
210        iter.into_iter()
211            .for_each(|(stat, value)| self.set(stat, value));
212    }
213}
214
215/// Merge helpers
216impl StatsSet {
217    /// Merge stats set `other` into `self`, with the semantic assumption that `other`
218    /// contains stats from a disjoint array that is *appended* to the array represented by `self`.
219    pub fn merge_ordered(mut self, other: &Self, dtype: &DType) -> Self {
220        self.as_mut_typed_ref(dtype)
221            .merge_ordered(&other.as_typed_ref(dtype));
222        self
223    }
224
225    /// Merge stats set `other` into `self`, from a disjoint array, with no ordering assumptions.
226    /// Stats that are not commutative (e.g., is_sorted) are dropped from the result.
227    pub fn merge_unordered(mut self, other: &Self, dtype: &DType) -> Self {
228        self.as_mut_typed_ref(dtype)
229            .merge_unordered(&other.as_typed_ref(dtype));
230        self
231    }
232
233    /// Given two sets of stats (of differing precision) for the same array, combine them
234    pub fn combine_sets(&mut self, other: &Self, dtype: &DType) -> VortexResult<()> {
235        self.as_mut_typed_ref(dtype)
236            .combine_sets(&other.as_typed_ref(dtype))
237    }
238}
239
240/// Builds the typed [`Scalar`] for `stat` over an array of `dtype` from its stored value.
241fn stat_scalar(values: &StatsSet, stat: Stat, dtype: &DType) -> Precision<Scalar> {
242    values.get_ref(stat).map(|value| {
243        let stat_dtype = stat
244            .dtype(dtype)
245            .vortex_expect("Must have valid dtype if value is present");
246        Scalar::validate(&stat_dtype, Some(value))
247            .vortex_expect("failed to construct a scalar statistic");
248
249        // SAFETY: The value was validated against the dtype above.
250        unsafe { Scalar::new_unchecked(stat_dtype, Some(value.clone())) }
251    })
252}
253
254pub struct TypedStatsSetRef<'a, 'b> {
255    pub values: &'a StatsSet,
256    pub dtype: &'b DType,
257}
258
259impl StatsProvider for TypedStatsSetRef<'_, '_> {
260    fn get(&self, stat: Stat) -> Precision<Scalar> {
261        stat_scalar(self.values, stat, self.dtype)
262    }
263
264    fn len(&self) -> usize {
265        self.values.len()
266    }
267}
268
269pub struct MutTypedStatsSetRef<'a, 'b> {
270    pub values: &'a mut StatsSet,
271    pub dtype: &'b DType,
272}
273
274impl MutTypedStatsSetRef<'_, '_> {
275    /// Set the stat `stat` to `value`.
276    pub fn set(&mut self, stat: Stat, value: Precision<ScalarValue>) {
277        self.values.set(stat, value);
278    }
279
280    /// Clear the stat `stat` from the set.
281    pub fn clear(&mut self, stat: Stat) {
282        self.values.clear(stat);
283    }
284}
285
286impl StatsProvider for MutTypedStatsSetRef<'_, '_> {
287    fn get(&self, stat: Stat) -> Precision<Scalar> {
288        stat_scalar(self.values, stat, self.dtype)
289    }
290
291    fn len(&self) -> usize {
292        self.values.len()
293    }
294}
295
296// Merge helpers
297impl MutTypedStatsSetRef<'_, '_> {
298    /// Merge stats set `other` into `self`, with the semantic assumption that `other`
299    /// contains stats from a disjoint array that is *appended* to the array represented by `self`.
300    pub fn merge_ordered(mut self, other: &TypedStatsSetRef) -> Self {
301        for s in all::<Stat>() {
302            match s {
303                Stat::IsConstant => self.merge_is_constant(other),
304                Stat::IsSorted => self.merge_is_sorted(other),
305                Stat::IsStrictSorted => self.merge_is_strict_sorted(other),
306                Stat::Max => self.merge_max(other),
307                Stat::Min => self.merge_min(other),
308                Stat::Sum => self.merge_sum(other),
309                Stat::NullCount => self.merge_null_count(other),
310                Stat::UncompressedSizeInBytes => self.merge_uncompressed_size_in_bytes(other),
311                Stat::NaNCount => self.merge_nan_count(other),
312            }
313        }
314
315        self
316    }
317
318    /// Merge stats set `other` into `self`, from a disjoint array, with no ordering assumptions.
319    /// Stats that are not commutative (e.g., is_sorted) are dropped from the result.
320    pub fn merge_unordered(mut self, other: &TypedStatsSetRef) -> Self {
321        for s in all::<Stat>() {
322            if !s.is_commutative() {
323                self.clear(s);
324                continue;
325            }
326
327            match s {
328                Stat::IsConstant => self.merge_is_constant(other),
329                Stat::Max => self.merge_max(other),
330                Stat::Min => self.merge_min(other),
331                Stat::Sum => self.merge_sum(other),
332                Stat::NullCount => self.merge_null_count(other),
333                Stat::UncompressedSizeInBytes => self.merge_uncompressed_size_in_bytes(other),
334                Stat::IsSorted | Stat::IsStrictSorted => {
335                    unreachable!("not commutative")
336                }
337                Stat::NaNCount => self.merge_nan_count(other),
338            }
339        }
340
341        self
342    }
343
344    /// Given two sets of stats (of differing precision) for the same array, combine them
345    pub fn combine_sets(&mut self, other: &TypedStatsSetRef) -> VortexResult<()> {
346        let other_stats: Vec<_> = other.values.iter().map(|(stat, _)| *stat).collect();
347        for s in other_stats {
348            match s {
349                Stat::Max => self.combine_bound::<Max>(other)?,
350                Stat::Min => self.combine_bound::<Min>(other)?,
351                Stat::UncompressedSizeInBytes => {
352                    self.combine_bound::<UncompressedSizeInBytes>(other)?
353                }
354                Stat::IsConstant => self.combine_bool_stat::<IsConstant>(other)?,
355                Stat::IsSorted => self.combine_bool_stat::<IsSorted>(other)?,
356                Stat::IsStrictSorted => self.combine_bool_stat::<IsStrictSorted>(other)?,
357                Stat::NullCount => self.combine_bound::<NullCount>(other)?,
358                Stat::Sum => self.combine_bound::<Sum>(other)?,
359                Stat::NaNCount => self.combine_bound::<NaNCount>(other)?,
360            }
361        }
362        Ok(())
363    }
364
365    fn combine_bound<S: StatType<Scalar>>(&mut self, other: &TypedStatsSetRef) -> VortexResult<()>
366    where
367        S::Bound: StatBound<Scalar> + Debug + Eq + PartialEq,
368    {
369        match (self.get_scalar_bound::<S>(), other.get_scalar_bound::<S>()) {
370            (Some(m1), Some(m2)) => {
371                let meet = m1
372                    .intersection(&m2)
373                    .vortex_expect("can always compare scalar")
374                    .ok_or_else(|| {
375                        vortex_err!("{:?} bounds ({m1:?}, {m2:?}) do not overlap", S::STAT)
376                    })?;
377                if meet != m1 {
378                    self.set(
379                        S::STAT,
380                        meet.into_value().map(|s| {
381                            s.into_value()
382                                .vortex_expect("stat scalar value cannot be null")
383                        }),
384                    );
385                }
386            }
387            (None, Some(m)) => self.set(
388                S::STAT,
389                m.into_value().map(|s| {
390                    s.into_value()
391                        .vortex_expect("stat scalar value cannot be null")
392                }),
393            ),
394            (Some(_), _) => (),
395            (None, None) => self.clear(S::STAT),
396        }
397        Ok(())
398    }
399
400    fn combine_bool_stat<S: StatType<bool>>(&mut self, other: &TypedStatsSetRef) -> VortexResult<()>
401    where
402        S::Bound: StatBound<bool> + Debug + Eq + PartialEq,
403    {
404        match (
405            self.get_as_bound::<S, bool>(),
406            other.get_as_bound::<S, bool>(),
407        ) {
408            (Some(m1), Some(m2)) => {
409                let intersection = m1
410                    .intersection(&m2)
411                    .vortex_expect("can always compare boolean")
412                    .ok_or_else(|| {
413                        vortex_err!("{:?} bounds ({m1:?}, {m2:?}) do not overlap", S::STAT)
414                    })?;
415                if intersection != m1 {
416                    self.set(S::STAT, intersection.into_value().map(ScalarValue::from));
417                }
418            }
419            (None, Some(m)) => self.set(S::STAT, m.into_value().map(ScalarValue::from)),
420            (Some(_), None) => (),
421            (None, None) => self.clear(S::STAT),
422        }
423        Ok(())
424    }
425
426    fn merge_min(&mut self, other: &TypedStatsSetRef) {
427        match (
428            self.get_scalar_bound::<Min>(),
429            other.get_scalar_bound::<Min>(),
430        ) {
431            (Some(m1), Some(m2)) => {
432                let meet = m1.union(&m2).vortex_expect("can compare scalar");
433                if meet != m1 {
434                    self.set(
435                        Stat::Min,
436                        meet.into_value().map(|s| {
437                            s.into_value()
438                                .vortex_expect("stat scalar value cannot be null")
439                        }),
440                    );
441                }
442            }
443            _ => self.clear(Stat::Min),
444        }
445    }
446
447    fn merge_max(&mut self, other: &TypedStatsSetRef) {
448        match (
449            self.get_scalar_bound::<Max>(),
450            other.get_scalar_bound::<Max>(),
451        ) {
452            (Some(m1), Some(m2)) => {
453                let meet = m1.union(&m2).vortex_expect("can compare scalar");
454                if meet != m1 {
455                    self.set(
456                        Stat::Max,
457                        meet.into_value().map(|s| {
458                            s.into_value()
459                                .vortex_expect("stat scalar value cannot be null")
460                        }),
461                    );
462                }
463            }
464            _ => self.clear(Stat::Max),
465        }
466    }
467
468    fn merge_sum(&mut self, other: &TypedStatsSetRef) {
469        match (
470            self.get_scalar_bound::<Sum>(),
471            other.get_scalar_bound::<Sum>(),
472        ) {
473            (Some(m1), Some(m2)) => {
474                // If the combined sum is exact, then we can sum them.
475                let merged = m1.zip(m2).as_exact().and_then(|(s1, s2)| match s1.dtype() {
476                    DType::Primitive(..) => s1
477                        .as_primitive()
478                        .checked_add(&s2.as_primitive())
479                        .and_then(|pscalar| pscalar.pvalue().map(ScalarValue::Primitive)),
480                    // Add widens the result precision, so the merged sum is only exact as a
481                    // stat if it still fits the summands' own decimal type.
482                    DType::Decimal(decimal_dtype, _) => s1
483                        .as_decimal()
484                        .checked_binary_numeric(
485                            &s2.as_decimal(),
486                            crate::scalar::NumericOperator::Add,
487                        )
488                        .ok()
489                        .flatten()
490                        .and_then(|scalar| scalar.as_decimal().decimal_value())
491                        .filter(|value| value.fits_in_precision(*decimal_dtype))
492                        .map(ScalarValue::Decimal),
493                    _ => None,
494                });
495
496                match merged {
497                    Some(scalar_value) => self.set(Stat::Sum, Precision::Exact(scalar_value)),
498                    // An overflow, an inexact bound or a non-numeric dtype leaves the combined
499                    // sum unknown. Keeping this set's own sum would report a partial total.
500                    None => self.clear(Stat::Sum),
501                }
502            }
503            _ => self.clear(Stat::Sum),
504        }
505    }
506
507    fn merge_is_constant(&mut self, other: &TypedStatsSetRef) {
508        let self_const = self.values.get_bool(Stat::IsConstant);
509        let other_const = other.values.get_bool(Stat::IsConstant);
510
511        match (self_const, other_const) {
512            // A non-constant part makes the whole array non-constant.
513            (Precision::Exact(false), _) | (_, Precision::Exact(false)) => {
514                self.set(Stat::IsConstant, Precision::exact(false));
515            }
516            (Precision::Exact(true), Precision::Exact(true)) => {
517                // Both mins have the same dtype, so comparing the values is the same as comparing
518                // scalars.
519                let self_min = self.values.get_ref(Stat::Min).as_exact();
520                let other_min = other.values.get_ref(Stat::Min).as_exact();
521
522                match (self_min, other_min) {
523                    (Some(self_min), Some(other_min)) => {
524                        let is_constant = self_min == other_min;
525                        self.set(Stat::IsConstant, Precision::exact(is_constant));
526                    }
527                    // Without both exact mins, e.g. for an all-null part, the parts cannot be
528                    // compared.
529                    _ => self.clear(Stat::IsConstant),
530                }
531            }
532            _ => self.clear(Stat::IsConstant),
533        }
534    }
535
536    fn merge_is_sorted(&mut self, other: &TypedStatsSetRef) {
537        self.merge_sortedness_stat(other, Stat::IsSorted, PartialOrd::le)
538    }
539
540    fn merge_is_strict_sorted(&mut self, other: &TypedStatsSetRef) {
541        self.merge_sortedness_stat(other, Stat::IsStrictSorted, PartialOrd::lt)
542    }
543
544    fn merge_sortedness_stat<F: Fn(&Scalar, &Scalar) -> bool>(
545        &mut self,
546        other: &TypedStatsSetRef,
547        stat: Stat,
548        cmp: F,
549    ) {
550        if (Precision::Exact(true), Precision::Exact(true))
551            == (self.values.get_bool(stat), other.values.get_bool(stat))
552        {
553            // There might be no stat because it was dropped, or it doesn't exist
554            // (e.g. an all null array).
555            // We assume that it was the dropped case since the doesn't exist might imply sorted,
556            // but this in-precision is correct.
557            if let (Some(self_max), Some(other_min)) = (
558                self.get_scalar_bound::<Max>().and_then(|v| v.max_value()),
559                other.get_scalar_bound::<Min>().and_then(|v| v.min_value()),
560            ) {
561                return if cmp(&self_max, &other_min) {
562                    // keep value
563                } else {
564                    self.set(stat, Precision::inexact(false));
565                };
566            }
567        }
568        self.clear(stat);
569    }
570
571    fn merge_null_count(&mut self, other: &TypedStatsSetRef) {
572        self.merge_sum_stat(Stat::NullCount, other)
573    }
574
575    fn merge_nan_count(&mut self, other: &TypedStatsSetRef) {
576        self.merge_sum_stat(Stat::NaNCount, other)
577    }
578
579    fn merge_uncompressed_size_in_bytes(&mut self, other: &TypedStatsSetRef) {
580        self.merge_sum_stat(Stat::UncompressedSizeInBytes, other)
581    }
582
583    fn merge_sum_stat(&mut self, stat: Stat, other: &TypedStatsSetRef) {
584        let merged = self
585            .values
586            .get_usize(stat)
587            .zip(other.values.get_usize(stat))
588            .map(|(l, r)| ScalarValue::from(l + r));
589
590        if merged.is_absent() {
591            self.clear(stat);
592        } else {
593            self.set(stat, merged);
594        }
595    }
596}
597
598#[cfg(test)]
599mod test {
600    use enum_iterator::all;
601    use itertools::Itertools;
602    use rstest::rstest;
603    use smallvec::smallvec;
604
605    use crate::VortexSessionExecute;
606    use crate::array_session;
607    use crate::arrays::PrimitiveArray;
608    use crate::dtype::DType;
609    use crate::dtype::DecimalDType;
610    use crate::dtype::MAX_PRECISION;
611    use crate::dtype::NativeDecimalType;
612    use crate::dtype::Nullability;
613    use crate::dtype::PType;
614    use crate::dtype::i256;
615    use crate::expr::stats::IsConstant;
616    use crate::expr::stats::Precision;
617    use crate::expr::stats::Stat;
618    use crate::expr::stats::StatsProvider;
619    use crate::expr::stats::StatsProviderExt;
620    use crate::scalar::DecimalValue;
621    use crate::scalar::ScalarValue;
622    use crate::stats::StatsSet;
623    use crate::stats::stats_set::Scalar;
624
625    #[test]
626    fn test_iter() {
627        // SAFETY: No duplicate stats.
628        let set = unsafe {
629            StatsSet::new_unchecked(smallvec![
630                (Stat::Max, Precision::exact(100)),
631                (Stat::Min, Precision::exact(42)),
632            ])
633        };
634        let mut iter = set.iter();
635        let first = iter.next().unwrap().clone();
636        assert_eq!(first.0, Stat::Max);
637        assert_eq!(
638            first.1.map(
639                |f| i32::try_from(&Scalar::try_new(PType::I32.into(), Some(f)).unwrap()).unwrap()
640            ),
641            Precision::exact(100)
642        );
643        let snd = iter.next().unwrap().clone();
644        assert_eq!(snd.0, Stat::Min);
645        assert_eq!(
646            snd.1.map(
647                |s| i32::try_from(&Scalar::try_new(PType::I32.into(), Some(s)).unwrap()).unwrap()
648            ),
649            Precision::exact(42)
650        );
651    }
652
653    #[test]
654    fn into_iter() {
655        // SAFETY: No duplicate stats.
656        let mut set = unsafe {
657            StatsSet::new_unchecked(smallvec![
658                (Stat::Max, Precision::exact(100)),
659                (Stat::Min, Precision::exact(42)),
660            ])
661        }
662        .into_iter();
663        let (stat, first) = set.next().unwrap();
664        assert_eq!(stat, Stat::Max);
665        assert_eq!(
666            first.map(
667                |f| i32::try_from(&Scalar::try_new(PType::I32.into(), Some(f)).unwrap()).unwrap()
668            ),
669            Precision::exact(100)
670        );
671        let snd = set.next().unwrap();
672        assert_eq!(snd.0, Stat::Min);
673        assert_eq!(
674            snd.1.map(
675                |s| i32::try_from(&Scalar::try_new(PType::I32.into(), Some(s)).unwrap()).unwrap()
676            ),
677            Precision::exact(42)
678        );
679    }
680
681    #[test]
682    fn merge_sums_overflow_clears() {
683        // A sum that cannot be combined exactly leaves the merged set with no Sum at all;
684        // retaining this set's own sum would report a partial total as the combined one.
685        let dtype = DType::Primitive(PType::I64, Nullability::NonNullable);
686        let merged = StatsSet::of(Stat::Sum, Precision::exact(i64::MAX))
687            .merge_ordered(&StatsSet::of(Stat::Sum, Precision::exact(i64::MAX)), &dtype);
688
689        assert!(merged.get(Stat::Sum).is_absent());
690    }
691
692    #[test]
693    fn merge_decimal_sums_out_of_precision_clears() {
694        // A decimal(70, 0) array's Sum stat is itself decimal(76, 0): `sum_decimal_dtype` adds
695        // ten digits of headroom but saturates at MAX_PRECISION, so at this input precision there
696        // is none left and two maximal sums cannot be combined exactly.
697        let dtype = DType::Decimal(DecimalDType::new(70, 0), Nullability::NonNullable);
698        let max = ScalarValue::Decimal(DecimalValue::I256(
699            <i256 as NativeDecimalType>::MAX_BY_PRECISION[usize::from(MAX_PRECISION)],
700        ));
701        let merged = StatsSet::of(Stat::Sum, Precision::exact(max.clone()))
702            .merge_ordered(&StatsSet::of(Stat::Sum, Precision::exact(max)), &dtype);
703
704        assert!(merged.get(Stat::Sum).is_absent());
705    }
706
707    #[test]
708    fn merge_constant() {
709        let first = StatsSet::from_iter([
710            (Stat::Min, Precision::exact(42)),
711            (Stat::IsConstant, Precision::exact(true)),
712        ])
713        .merge_ordered(
714            &StatsSet::from_iter([
715                (Stat::Min, Precision::inexact(42)),
716                (Stat::IsConstant, Precision::exact(true)),
717            ]),
718            &DType::Primitive(PType::I32, Nullability::NonNullable),
719        );
720
721        // With an inexact min, the two constant parts cannot be compared.
722        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
723        assert_eq!(
724            first_ref.get_as::<bool>(Stat::IsConstant),
725            Precision::Absent
726        );
727        assert_eq!(first_ref.get_as::<i32>(Stat::Min), Precision::exact(42));
728    }
729
730    #[rstest]
731    #[case::equal_mins(
732        Precision::exact(true),
733        42,
734        Precision::exact(true),
735        42,
736        Precision::exact(true)
737    )]
738    #[case::different_mins(
739        Precision::exact(true),
740        42,
741        Precision::exact(true),
742        7,
743        Precision::exact(false)
744    )]
745    #[case::first_not_constant(
746        Precision::exact(false),
747        42,
748        Precision::exact(true),
749        42,
750        Precision::exact(false)
751    )]
752    #[case::second_not_constant(
753        Precision::exact(true),
754        42,
755        Precision::exact(false),
756        42,
757        Precision::exact(false)
758    )]
759    #[case::inexact_constant(
760        Precision::inexact(true),
761        42,
762        Precision::exact(true),
763        42,
764        Precision::Absent
765    )]
766    fn merge_constant_parts(
767        #[case] first_const: Precision<bool>,
768        #[case] first_min: i32,
769        #[case] second_const: Precision<bool>,
770        #[case] second_min: i32,
771        #[case] expected: Precision<bool>,
772    ) {
773        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
774        let stats = |is_constant: Precision<bool>, min: i32| {
775            StatsSet::from_iter([
776                (Stat::Min, Precision::exact(min)),
777                (Stat::IsConstant, is_constant.map(ScalarValue::from)),
778            ])
779        };
780
781        let merged =
782            stats(first_const, first_min).merge_ordered(&stats(second_const, second_min), &dtype);
783
784        assert_eq!(
785            merged.as_typed_ref(&dtype).get_as::<bool>(Stat::IsConstant),
786            expected
787        );
788    }
789
790    #[test]
791    fn merge_into_min() {
792        let first = StatsSet::of(Stat::Min, Precision::exact(42)).merge_ordered(
793            &StatsSet::default(),
794            &DType::Primitive(PType::I32, Nullability::NonNullable),
795        );
796
797        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
798        assert!(first_ref.get(Stat::Min).is_absent());
799    }
800
801    #[test]
802    fn merge_from_min() {
803        let first = StatsSet::default().merge_ordered(
804            &StatsSet::of(Stat::Min, Precision::exact(42)),
805            &DType::Primitive(PType::I32, Nullability::NonNullable),
806        );
807
808        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
809        assert!(first_ref.get(Stat::Min).is_absent());
810    }
811
812    #[test]
813    fn merge_mins() {
814        let first = StatsSet::of(Stat::Min, Precision::exact(37)).merge_ordered(
815            &StatsSet::of(Stat::Min, Precision::exact(42)),
816            &DType::Primitive(PType::I32, Nullability::NonNullable),
817        );
818
819        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
820        assert_eq!(first_ref.get_as::<i32>(Stat::Min), Precision::exact(37));
821    }
822
823    #[test]
824    fn merge_into_bound_max() {
825        let first = StatsSet::of(Stat::Max, Precision::exact(42)).merge_ordered(
826            &StatsSet::default(),
827            &DType::Primitive(PType::I32, Nullability::NonNullable),
828        );
829        assert!(first.get(Stat::Max).is_absent());
830    }
831
832    #[test]
833    fn merge_from_max() {
834        let first = StatsSet::default().merge_ordered(
835            &StatsSet::of(Stat::Max, Precision::exact(42)),
836            &DType::Primitive(PType::I32, Nullability::NonNullable),
837        );
838        assert!(first.get(Stat::Max).is_absent());
839    }
840
841    #[test]
842    fn merge_maxes() {
843        let first = StatsSet::of(Stat::Max, Precision::exact(37)).merge_ordered(
844            &StatsSet::of(Stat::Max, Precision::exact(42)),
845            &DType::Primitive(PType::I32, Nullability::NonNullable),
846        );
847        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
848        assert_eq!(first_ref.get_as::<i32>(Stat::Max), Precision::exact(42));
849    }
850
851    #[test]
852    fn merge_maxes_bound() {
853        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
854        let first = StatsSet::of(Stat::Max, Precision::exact(42i32))
855            .merge_ordered(&StatsSet::of(Stat::Max, Precision::inexact(43i32)), &dtype);
856        let first_ref = first.as_typed_ref(&dtype);
857        assert_eq!(first_ref.get_as::<i32>(Stat::Max), Precision::inexact(43));
858    }
859
860    #[test]
861    fn merge_into_scalar() {
862        // Sum stats for primitive types are always the 64-bit version (i64 for signed, u64
863        // for unsigned, f64 for floats).
864        let first = StatsSet::of(Stat::Sum, Precision::exact(42i64)).merge_ordered(
865            &StatsSet::default(),
866            &DType::Primitive(PType::I32, Nullability::NonNullable),
867        );
868        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
869        assert!(first_ref.get(Stat::Sum).is_absent());
870    }
871
872    #[test]
873    fn merge_from_scalar() {
874        // Sum stats for primitive types are always the 64-bit version (i64 for signed, u64
875        // for unsigned, f64 for floats).
876        let first = StatsSet::default().merge_ordered(
877            &StatsSet::of(Stat::Sum, Precision::exact(42i64)),
878            &DType::Primitive(PType::I32, Nullability::NonNullable),
879        );
880        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
881        assert!(first_ref.get(Stat::Sum).is_absent());
882    }
883
884    #[test]
885    fn merge_scalars() {
886        // Sum stats for primitive types are always the 64-bit version (i64 for signed, u64
887        // for unsigned, f64 for floats).
888        let first = StatsSet::of(Stat::Sum, Precision::exact(37i64)).merge_ordered(
889            &StatsSet::of(Stat::Sum, Precision::exact(42i64)),
890            &DType::Primitive(PType::I32, Nullability::NonNullable),
891        );
892        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
893        assert_eq!(first_ref.get_as::<i64>(Stat::Sum), Precision::exact(79i64));
894    }
895
896    #[test]
897    fn merge_into_sortedness() {
898        let first = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true)).merge_ordered(
899            &StatsSet::default(),
900            &DType::Primitive(PType::I32, Nullability::NonNullable),
901        );
902        assert!(first.get(Stat::IsStrictSorted).is_absent());
903    }
904
905    #[test]
906    fn merge_from_sortedness() {
907        let first = StatsSet::default().merge_ordered(
908            &StatsSet::of(Stat::IsStrictSorted, Precision::exact(true)),
909            &DType::Primitive(PType::I32, Nullability::NonNullable),
910        );
911        assert!(first.get(Stat::IsStrictSorted).is_absent());
912    }
913
914    #[test]
915    fn merge_sortedness() {
916        let mut first = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true));
917        first.set(Stat::Max, Precision::exact(1));
918        let mut second = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true));
919        second.set(Stat::Min, Precision::exact(2));
920        first = first.merge_ordered(
921            &second,
922            &DType::Primitive(PType::I32, Nullability::NonNullable),
923        );
924
925        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
926        assert_eq!(
927            first_ref.get_as::<bool>(Stat::IsStrictSorted),
928            Precision::exact(true)
929        );
930    }
931
932    #[test]
933    fn merge_sortedness_out_of_order() {
934        let mut first = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true));
935        first.set(Stat::Min, Precision::exact(1));
936        let mut second = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true));
937        second.set(Stat::Max, Precision::exact(2));
938        second = second.merge_ordered(
939            &first,
940            &DType::Primitive(PType::I32, Nullability::NonNullable),
941        );
942
943        let second_ref =
944            second.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
945        assert_eq!(
946            second_ref.get_as::<bool>(Stat::IsStrictSorted),
947            Precision::inexact(false)
948        );
949    }
950
951    #[test]
952    fn merge_sortedness_only_one_sorted() {
953        let mut first = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true));
954        first.set(Stat::Max, Precision::exact(1));
955        let mut second = StatsSet::of(Stat::IsStrictSorted, Precision::exact(false));
956        second.set(Stat::Min, Precision::exact(2));
957        first.merge_ordered(
958            &second,
959            &DType::Primitive(PType::I32, Nullability::NonNullable),
960        );
961
962        let second_ref =
963            second.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
964        assert_eq!(
965            second_ref.get_as::<bool>(Stat::IsStrictSorted),
966            Precision::exact(false)
967        );
968    }
969
970    #[test]
971    fn merge_sortedness_missing_min() {
972        let mut first = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true));
973        first.set(Stat::Max, Precision::exact(1));
974        let second = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true));
975        first = first.merge_ordered(
976            &second,
977            &DType::Primitive(PType::I32, Nullability::NonNullable),
978        );
979        assert!(first.get(Stat::IsStrictSorted).is_absent());
980    }
981
982    #[test]
983    fn merge_sortedness_bound_min() {
984        let mut first = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true));
985        first.set(Stat::Max, Precision::exact(1));
986        let mut second = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true));
987        second.set(Stat::Min, Precision::inexact(2));
988        first = first.merge_ordered(
989            &second,
990            &DType::Primitive(PType::I32, Nullability::NonNullable),
991        );
992
993        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
994        assert_eq!(
995            first_ref.get_as::<bool>(Stat::IsStrictSorted),
996            Precision::exact(true)
997        );
998    }
999
1000    #[test]
1001    fn merge_unordered() {
1002        let array =
1003            PrimitiveArray::from_option_iter([Some(1), None, Some(2), Some(42), Some(10000), None]);
1004        let all_stats = all::<Stat>()
1005            .filter(|s| !matches!(s, Stat::Sum))
1006            .filter(|s| !matches!(s, Stat::NaNCount))
1007            .collect_vec();
1008        array
1009            .statistics()
1010            .compute_all(&all_stats, &mut array_session().create_execution_ctx())
1011            .unwrap();
1012
1013        let stats = array.statistics().to_owned();
1014        for stat in &all_stats {
1015            assert!(!stats.get(*stat).is_absent(), "Stat {stat} is missing");
1016        }
1017
1018        let merged = stats.clone().merge_unordered(
1019            &stats,
1020            &DType::Primitive(PType::I32, Nullability::NonNullable),
1021        );
1022        for stat in &all_stats {
1023            assert_eq!(
1024                !merged.get(*stat).is_absent(),
1025                stat.is_commutative(),
1026                "Stat {stat} remains after merge_unordered despite not being commutative, or was removed despite being commutative"
1027            )
1028        }
1029
1030        let merged_ref = merged.as_typed_ref(&DType::Primitive(PType::I32, Nullability::Nullable));
1031        let stats_ref = stats.as_typed_ref(&DType::Primitive(PType::I32, Nullability::Nullable));
1032
1033        assert_eq!(
1034            merged_ref.get_as::<i32>(Stat::Min),
1035            stats_ref.get_as::<i32>(Stat::Min)
1036        );
1037        assert_eq!(
1038            merged_ref.get_as::<i32>(Stat::Max),
1039            stats_ref.get_as::<i32>(Stat::Max)
1040        );
1041        assert_eq!(
1042            merged_ref.get_as::<u64>(Stat::NullCount),
1043            stats_ref.get_as::<u64>(Stat::NullCount).map(|s| s * 2)
1044        );
1045    }
1046
1047    #[test]
1048    fn merge_min_bound_same() {
1049        // Merging a stat with a bound and another with an exact results in exact stat.
1050        // since bound for min is a lower bound, it can in fact contain any value >= bound.
1051        let merged = StatsSet::of(Stat::Min, Precision::inexact(5)).merge_ordered(
1052            &StatsSet::of(Stat::Min, Precision::exact(5)),
1053            &DType::Primitive(PType::I32, Nullability::NonNullable),
1054        );
1055        let merged_ref =
1056            merged.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
1057        assert_eq!(merged_ref.get_as::<i32>(Stat::Min), Precision::exact(5));
1058    }
1059
1060    #[test]
1061    fn merge_min_bound_bound_lower() {
1062        let merged = StatsSet::of(Stat::Min, Precision::inexact(4)).merge_ordered(
1063            &StatsSet::of(Stat::Min, Precision::exact(5)),
1064            &DType::Primitive(PType::I32, Nullability::NonNullable),
1065        );
1066        let merged_ref =
1067            merged.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
1068        assert_eq!(merged_ref.get_as::<i32>(Stat::Min), Precision::inexact(4));
1069    }
1070
1071    #[test]
1072    fn test_combine_is_constant() {
1073        {
1074            let mut stats = StatsSet::of(Stat::IsConstant, Precision::exact(true));
1075            let stats2 = StatsSet::of(Stat::IsConstant, Precision::exact(true));
1076            let mut stats_ref =
1077                stats.as_mut_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
1078            stats_ref
1079                .combine_bool_stat::<IsConstant>(
1080                    &stats2.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable)),
1081                )
1082                .unwrap();
1083            assert_eq!(
1084                stats_ref.get_as::<bool>(Stat::IsConstant),
1085                Precision::exact(true)
1086            );
1087        }
1088
1089        {
1090            let mut stats = StatsSet::of(Stat::IsConstant, Precision::exact(true));
1091            let stats2 = StatsSet::of(Stat::IsConstant, Precision::inexact(false));
1092            let mut stats_ref =
1093                stats.as_mut_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
1094            stats_ref
1095                .combine_bool_stat::<IsConstant>(
1096                    &stats2.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable)),
1097                )
1098                .unwrap();
1099            assert_eq!(
1100                stats_ref.get_as::<bool>(Stat::IsConstant),
1101                Precision::exact(true)
1102            );
1103        }
1104
1105        {
1106            let mut stats = StatsSet::of(Stat::IsConstant, Precision::exact(false));
1107            let stats2 = StatsSet::of(Stat::IsConstant, Precision::inexact(false));
1108            let mut stats_ref =
1109                stats.as_mut_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
1110            stats_ref
1111                .combine_bool_stat::<IsConstant>(
1112                    &stats2.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable)),
1113                )
1114                .unwrap();
1115            assert_eq!(
1116                stats_ref.get_as::<bool>(Stat::IsConstant),
1117                Precision::exact(false)
1118            );
1119        }
1120    }
1121
1122    #[test]
1123    fn test_combine_sets_boolean_conflict() {
1124        let mut stats1 = StatsSet::from_iter([
1125            (Stat::IsConstant, Precision::exact(true)),
1126            (Stat::IsSorted, Precision::exact(true)),
1127        ]);
1128
1129        let stats2 = StatsSet::from_iter([
1130            (Stat::IsConstant, Precision::exact(false)),
1131            (Stat::IsSorted, Precision::exact(true)),
1132        ]);
1133
1134        let result = stats1.combine_sets(
1135            &stats2,
1136            &DType::Primitive(PType::I32, Nullability::NonNullable),
1137        );
1138        assert!(result.is_err());
1139    }
1140
1141    #[test]
1142    fn test_combine_sets_with_missing_stats() {
1143        let mut stats1 = StatsSet::from_iter([
1144            (Stat::Min, Precision::exact(42)),
1145            (Stat::UncompressedSizeInBytes, Precision::exact(1000)),
1146        ]);
1147
1148        let stats2 = StatsSet::from_iter([
1149            (Stat::Max, Precision::exact(100)),
1150            (Stat::IsStrictSorted, Precision::exact(true)),
1151        ]);
1152
1153        stats1
1154            .combine_sets(
1155                &stats2,
1156                &DType::Primitive(PType::I32, Nullability::NonNullable),
1157            )
1158            .unwrap();
1159
1160        let stats_ref =
1161            stats1.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
1162
1163        // Min should remain unchanged
1164        assert_eq!(stats_ref.get_as::<i32>(Stat::Min), Precision::exact(42));
1165        // Max should be added
1166        assert_eq!(stats_ref.get_as::<i32>(Stat::Max), Precision::exact(100));
1167        // IsStrictSorted should be added
1168        assert_eq!(
1169            stats_ref.get_as::<bool>(Stat::IsStrictSorted),
1170            Precision::exact(true)
1171        );
1172    }
1173
1174    #[test]
1175    fn test_combine_sets_with_inexact() {
1176        let mut stats1 = StatsSet::from_iter([
1177            (Stat::Min, Precision::exact(42)),
1178            (Stat::Max, Precision::inexact(100)),
1179            (Stat::IsConstant, Precision::exact(false)),
1180        ]);
1181
1182        let stats2 = StatsSet::from_iter([
1183            // Must ensure Min from stats2 is <= Min from stats1
1184            (Stat::Min, Precision::inexact(40)),
1185            (Stat::Max, Precision::exact(90)),
1186            (Stat::IsSorted, Precision::exact(true)),
1187        ]);
1188
1189        stats1
1190            .combine_sets(
1191                &stats2,
1192                &DType::Primitive(PType::I32, Nullability::NonNullable),
1193            )
1194            .unwrap();
1195
1196        let stats_ref =
1197            stats1.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
1198
1199        // Min should remain unchanged since it's more restrictive than the inexact value
1200        assert_eq!(stats_ref.get_as::<i32>(Stat::Min), Precision::exact(42));
1201        // Check that max was updated with the exact value
1202        assert_eq!(stats_ref.get_as::<i32>(Stat::Max), Precision::exact(90));
1203        // Check that IsSorted was added
1204        assert_eq!(
1205            stats_ref.get_as::<bool>(Stat::IsSorted),
1206            Precision::exact(true)
1207        );
1208    }
1209}
vortex_array/stats/stats_set.rs

vortex_array/stats/
stats_set.rs
