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

vortex_array/stats/
stats_set.rs
