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 combine sum is exact, then we can sum them.
447                if let Some(scalar_value) =
448                    m1.zip(m2).as_exact().and_then(|(s1, s2)| match s1.dtype() {
449                        DType::Primitive(..) => s1
450                            .as_primitive()
451                            .checked_add(&s2.as_primitive())
452                            .and_then(|pscalar| pscalar.pvalue().map(ScalarValue::Primitive)),
453                        DType::Decimal(..) => s1
454                            .as_decimal()
455                            .checked_binary_numeric(
456                                &s2.as_decimal(),
457                                crate::scalar::NumericOperator::Add,
458                            )
459                            .map(|scalar| {
460                                ScalarValue::Decimal(
461                                    scalar
462                                        .decimal_value()
463                                        .vortex_expect("no decimal value in scalar"),
464                                )
465                            }),
466                        _ => None,
467                    })
468                {
469                    self.set(Stat::Sum, Precision::Exact(scalar_value));
470                }
471            }
472            _ => self.clear(Stat::Sum),
473        }
474    }
475
476    fn merge_is_constant(&mut self, other: &TypedStatsSetRef) {
477        let self_const = self.get_as(Stat::IsConstant);
478        let other_const = other.get_as(Stat::IsConstant);
479        let self_min = self.get(Stat::Min);
480        let other_min = other.get(Stat::Min);
481
482        if let (Some(self_const), Some(other_const), Some(self_min), Some(other_min)) = (
483            self_const.as_exact(),
484            other_const.as_exact(),
485            self_min.as_exact(),
486            other_min.as_exact(),
487        ) {
488            if self_const && other_const && self_min == other_min {
489                self.set(Stat::IsConstant, Precision::exact(true));
490            } else {
491                self.set(Stat::IsConstant, Precision::inexact(false));
492            }
493        }
494        self.set(Stat::IsConstant, Precision::exact(false));
495    }
496
497    fn merge_is_sorted(&mut self, other: &TypedStatsSetRef) {
498        self.merge_sortedness_stat(other, Stat::IsSorted, PartialOrd::le)
499    }
500
501    fn merge_is_strict_sorted(&mut self, other: &TypedStatsSetRef) {
502        self.merge_sortedness_stat(other, Stat::IsStrictSorted, PartialOrd::lt)
503    }
504
505    fn merge_sortedness_stat<F: Fn(&Scalar, &Scalar) -> bool>(
506        &mut self,
507        other: &TypedStatsSetRef,
508        stat: Stat,
509        cmp: F,
510    ) {
511        if (Precision::Exact(true), Precision::Exact(true))
512            == (self.get_as(stat), other.get_as(stat))
513        {
514            // There might be no stat because it was dropped, or it doesn't exist
515            // (e.g. an all null array).
516            // We assume that it was the dropped case since the doesn't exist might imply sorted,
517            // but this in-precision is correct.
518            if let (Some(self_max), Some(other_min)) = (
519                self.get_scalar_bound::<Max>().and_then(|v| v.max_value()),
520                other.get_scalar_bound::<Min>().and_then(|v| v.min_value()),
521            ) {
522                return if cmp(&self_max, &other_min) {
523                    // keep value
524                } else {
525                    self.set(stat, Precision::inexact(false));
526                };
527            }
528        }
529        self.clear(stat);
530    }
531
532    fn merge_null_count(&mut self, other: &TypedStatsSetRef) {
533        self.merge_sum_stat(Stat::NullCount, other)
534    }
535
536    fn merge_nan_count(&mut self, other: &TypedStatsSetRef) {
537        self.merge_sum_stat(Stat::NaNCount, other)
538    }
539
540    fn merge_uncompressed_size_in_bytes(&mut self, other: &TypedStatsSetRef) {
541        self.merge_sum_stat(Stat::UncompressedSizeInBytes, other)
542    }
543
544    fn merge_sum_stat(&mut self, stat: Stat, other: &TypedStatsSetRef) {
545        let merged = self
546            .get_as::<usize>(stat)
547            .zip(other.get_as::<usize>(stat))
548            .map(|(l, r)| ScalarValue::from(l + r));
549
550        if merged.is_absent() {
551            self.clear(stat);
552        } else {
553            self.set(stat, merged);
554        }
555    }
556}
557
558#[cfg(test)]
559mod test {
560    use enum_iterator::all;
561    use itertools::Itertools;
562    use smallvec::smallvec;
563
564    use crate::VortexSessionExecute;
565    use crate::array_session;
566    use crate::arrays::PrimitiveArray;
567    use crate::dtype::DType;
568    use crate::dtype::Nullability;
569    use crate::dtype::PType;
570    use crate::expr::stats::IsConstant;
571    use crate::expr::stats::Precision;
572    use crate::expr::stats::Stat;
573    use crate::expr::stats::StatsProvider;
574    use crate::expr::stats::StatsProviderExt;
575    use crate::stats::StatsSet;
576    use crate::stats::stats_set::Scalar;
577
578    #[test]
579    fn test_iter() {
580        // SAFETY: No duplicate stats.
581        let set = unsafe {
582            StatsSet::new_unchecked(smallvec![
583                (Stat::Max, Precision::exact(100)),
584                (Stat::Min, Precision::exact(42)),
585            ])
586        };
587        let mut iter = set.iter();
588        let first = iter.next().unwrap().clone();
589        assert_eq!(first.0, Stat::Max);
590        assert_eq!(
591            first.1.map(
592                |f| i32::try_from(&Scalar::try_new(PType::I32.into(), Some(f)).unwrap()).unwrap()
593            ),
594            Precision::exact(100)
595        );
596        let snd = iter.next().unwrap().clone();
597        assert_eq!(snd.0, Stat::Min);
598        assert_eq!(
599            snd.1.map(
600                |s| i32::try_from(&Scalar::try_new(PType::I32.into(), Some(s)).unwrap()).unwrap()
601            ),
602            Precision::exact(42)
603        );
604    }
605
606    #[test]
607    fn into_iter() {
608        // SAFETY: No duplicate stats.
609        let mut set = unsafe {
610            StatsSet::new_unchecked(smallvec![
611                (Stat::Max, Precision::exact(100)),
612                (Stat::Min, Precision::exact(42)),
613            ])
614        }
615        .into_iter();
616        let (stat, first) = set.next().unwrap();
617        assert_eq!(stat, Stat::Max);
618        assert_eq!(
619            first.map(
620                |f| i32::try_from(&Scalar::try_new(PType::I32.into(), Some(f)).unwrap()).unwrap()
621            ),
622            Precision::exact(100)
623        );
624        let snd = set.next().unwrap();
625        assert_eq!(snd.0, Stat::Min);
626        assert_eq!(
627            snd.1.map(
628                |s| i32::try_from(&Scalar::try_new(PType::I32.into(), Some(s)).unwrap()).unwrap()
629            ),
630            Precision::exact(42)
631        );
632    }
633
634    #[test]
635    fn merge_constant() {
636        let first = StatsSet::from_iter([
637            (Stat::Min, Precision::exact(42)),
638            (Stat::IsConstant, Precision::exact(true)),
639        ])
640        .merge_ordered(
641            &StatsSet::from_iter([
642                (Stat::Min, Precision::inexact(42)),
643                (Stat::IsConstant, Precision::exact(true)),
644            ]),
645            &DType::Primitive(PType::I32, Nullability::NonNullable),
646        );
647
648        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
649        assert_eq!(
650            first_ref.get_as::<bool>(Stat::IsConstant),
651            Precision::exact(false)
652        );
653        assert_eq!(first_ref.get_as::<i32>(Stat::Min), Precision::exact(42));
654    }
655
656    #[test]
657    fn merge_into_min() {
658        let first = StatsSet::of(Stat::Min, Precision::exact(42)).merge_ordered(
659            &StatsSet::default(),
660            &DType::Primitive(PType::I32, Nullability::NonNullable),
661        );
662
663        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
664        assert!(first_ref.get(Stat::Min).is_absent());
665    }
666
667    #[test]
668    fn merge_from_min() {
669        let first = StatsSet::default().merge_ordered(
670            &StatsSet::of(Stat::Min, Precision::exact(42)),
671            &DType::Primitive(PType::I32, Nullability::NonNullable),
672        );
673
674        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
675        assert!(first_ref.get(Stat::Min).is_absent());
676    }
677
678    #[test]
679    fn merge_mins() {
680        let first = StatsSet::of(Stat::Min, Precision::exact(37)).merge_ordered(
681            &StatsSet::of(Stat::Min, Precision::exact(42)),
682            &DType::Primitive(PType::I32, Nullability::NonNullable),
683        );
684
685        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
686        assert_eq!(first_ref.get_as::<i32>(Stat::Min), Precision::exact(37));
687    }
688
689    #[test]
690    fn merge_into_bound_max() {
691        let first = StatsSet::of(Stat::Max, Precision::exact(42)).merge_ordered(
692            &StatsSet::default(),
693            &DType::Primitive(PType::I32, Nullability::NonNullable),
694        );
695        assert!(first.get(Stat::Max).is_absent());
696    }
697
698    #[test]
699    fn merge_from_max() {
700        let first = StatsSet::default().merge_ordered(
701            &StatsSet::of(Stat::Max, Precision::exact(42)),
702            &DType::Primitive(PType::I32, Nullability::NonNullable),
703        );
704        assert!(first.get(Stat::Max).is_absent());
705    }
706
707    #[test]
708    fn merge_maxes() {
709        let first = StatsSet::of(Stat::Max, Precision::exact(37)).merge_ordered(
710            &StatsSet::of(Stat::Max, Precision::exact(42)),
711            &DType::Primitive(PType::I32, Nullability::NonNullable),
712        );
713        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
714        assert_eq!(first_ref.get_as::<i32>(Stat::Max), Precision::exact(42));
715    }
716
717    #[test]
718    fn merge_maxes_bound() {
719        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
720        let first = StatsSet::of(Stat::Max, Precision::exact(42i32))
721            .merge_ordered(&StatsSet::of(Stat::Max, Precision::inexact(43i32)), &dtype);
722        let first_ref = first.as_typed_ref(&dtype);
723        assert_eq!(first_ref.get_as::<i32>(Stat::Max), Precision::inexact(43));
724    }
725
726    #[test]
727    fn merge_into_scalar() {
728        // Sum stats for primitive types are always the 64-bit version (i64 for signed, u64
729        // for unsigned, f64 for floats).
730        let first = StatsSet::of(Stat::Sum, Precision::exact(42i64)).merge_ordered(
731            &StatsSet::default(),
732            &DType::Primitive(PType::I32, Nullability::NonNullable),
733        );
734        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
735        assert!(first_ref.get(Stat::Sum).is_absent());
736    }
737
738    #[test]
739    fn merge_from_scalar() {
740        // Sum stats for primitive types are always the 64-bit version (i64 for signed, u64
741        // for unsigned, f64 for floats).
742        let first = StatsSet::default().merge_ordered(
743            &StatsSet::of(Stat::Sum, Precision::exact(42i64)),
744            &DType::Primitive(PType::I32, Nullability::NonNullable),
745        );
746        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
747        assert!(first_ref.get(Stat::Sum).is_absent());
748    }
749
750    #[test]
751    fn merge_scalars() {
752        // Sum stats for primitive types are always the 64-bit version (i64 for signed, u64
753        // for unsigned, f64 for floats).
754        let first = StatsSet::of(Stat::Sum, Precision::exact(37i64)).merge_ordered(
755            &StatsSet::of(Stat::Sum, Precision::exact(42i64)),
756            &DType::Primitive(PType::I32, Nullability::NonNullable),
757        );
758        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
759        assert_eq!(first_ref.get_as::<i64>(Stat::Sum), Precision::exact(79i64));
760    }
761
762    #[test]
763    fn merge_into_sortedness() {
764        let first = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true)).merge_ordered(
765            &StatsSet::default(),
766            &DType::Primitive(PType::I32, Nullability::NonNullable),
767        );
768        assert!(first.get(Stat::IsStrictSorted).is_absent());
769    }
770
771    #[test]
772    fn merge_from_sortedness() {
773        let first = StatsSet::default().merge_ordered(
774            &StatsSet::of(Stat::IsStrictSorted, Precision::exact(true)),
775            &DType::Primitive(PType::I32, Nullability::NonNullable),
776        );
777        assert!(first.get(Stat::IsStrictSorted).is_absent());
778    }
779
780    #[test]
781    fn merge_sortedness() {
782        let mut first = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true));
783        first.set(Stat::Max, Precision::exact(1));
784        let mut second = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true));
785        second.set(Stat::Min, Precision::exact(2));
786        first = first.merge_ordered(
787            &second,
788            &DType::Primitive(PType::I32, Nullability::NonNullable),
789        );
790
791        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
792        assert_eq!(
793            first_ref.get_as::<bool>(Stat::IsStrictSorted),
794            Precision::exact(true)
795        );
796    }
797
798    #[test]
799    fn merge_sortedness_out_of_order() {
800        let mut first = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true));
801        first.set(Stat::Min, Precision::exact(1));
802        let mut second = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true));
803        second.set(Stat::Max, Precision::exact(2));
804        second = second.merge_ordered(
805            &first,
806            &DType::Primitive(PType::I32, Nullability::NonNullable),
807        );
808
809        let second_ref =
810            second.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
811        assert_eq!(
812            second_ref.get_as::<bool>(Stat::IsStrictSorted),
813            Precision::inexact(false)
814        );
815    }
816
817    #[test]
818    fn merge_sortedness_only_one_sorted() {
819        let mut first = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true));
820        first.set(Stat::Max, Precision::exact(1));
821        let mut second = StatsSet::of(Stat::IsStrictSorted, Precision::exact(false));
822        second.set(Stat::Min, Precision::exact(2));
823        first.merge_ordered(
824            &second,
825            &DType::Primitive(PType::I32, Nullability::NonNullable),
826        );
827
828        let second_ref =
829            second.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
830        assert_eq!(
831            second_ref.get_as::<bool>(Stat::IsStrictSorted),
832            Precision::exact(false)
833        );
834    }
835
836    #[test]
837    fn merge_sortedness_missing_min() {
838        let mut first = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true));
839        first.set(Stat::Max, Precision::exact(1));
840        let second = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true));
841        first = first.merge_ordered(
842            &second,
843            &DType::Primitive(PType::I32, Nullability::NonNullable),
844        );
845        assert!(first.get(Stat::IsStrictSorted).is_absent());
846    }
847
848    #[test]
849    fn merge_sortedness_bound_min() {
850        let mut first = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true));
851        first.set(Stat::Max, Precision::exact(1));
852        let mut second = StatsSet::of(Stat::IsStrictSorted, Precision::exact(true));
853        second.set(Stat::Min, Precision::inexact(2));
854        first = first.merge_ordered(
855            &second,
856            &DType::Primitive(PType::I32, Nullability::NonNullable),
857        );
858
859        let first_ref = first.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
860        assert_eq!(
861            first_ref.get_as::<bool>(Stat::IsStrictSorted),
862            Precision::exact(true)
863        );
864    }
865
866    #[test]
867    fn merge_unordered() {
868        let array =
869            PrimitiveArray::from_option_iter([Some(1), None, Some(2), Some(42), Some(10000), None]);
870        let all_stats = all::<Stat>()
871            .filter(|s| !matches!(s, Stat::Sum))
872            .filter(|s| !matches!(s, Stat::NaNCount))
873            .collect_vec();
874        array
875            .statistics()
876            .compute_all(&all_stats, &mut array_session().create_execution_ctx())
877            .unwrap();
878
879        let stats = array.statistics().to_owned();
880        for stat in &all_stats {
881            assert!(!stats.get(*stat).is_absent(), "Stat {stat} is missing");
882        }
883
884        let merged = stats.clone().merge_unordered(
885            &stats,
886            &DType::Primitive(PType::I32, Nullability::NonNullable),
887        );
888        for stat in &all_stats {
889            assert_eq!(
890                !merged.get(*stat).is_absent(),
891                stat.is_commutative(),
892                "Stat {stat} remains after merge_unordered despite not being commutative, or was removed despite being commutative"
893            )
894        }
895
896        let merged_ref = merged.as_typed_ref(&DType::Primitive(PType::I32, Nullability::Nullable));
897        let stats_ref = stats.as_typed_ref(&DType::Primitive(PType::I32, Nullability::Nullable));
898
899        assert_eq!(
900            merged_ref.get_as::<i32>(Stat::Min),
901            stats_ref.get_as::<i32>(Stat::Min)
902        );
903        assert_eq!(
904            merged_ref.get_as::<i32>(Stat::Max),
905            stats_ref.get_as::<i32>(Stat::Max)
906        );
907        assert_eq!(
908            merged_ref.get_as::<u64>(Stat::NullCount),
909            stats_ref.get_as::<u64>(Stat::NullCount).map(|s| s * 2)
910        );
911    }
912
913    #[test]
914    fn merge_min_bound_same() {
915        // Merging a stat with a bound and another with an exact results in exact stat.
916        // since bound for min is a lower bound, it can in fact contain any value >= bound.
917        let merged = StatsSet::of(Stat::Min, Precision::inexact(5)).merge_ordered(
918            &StatsSet::of(Stat::Min, Precision::exact(5)),
919            &DType::Primitive(PType::I32, Nullability::NonNullable),
920        );
921        let merged_ref =
922            merged.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
923        assert_eq!(merged_ref.get_as::<i32>(Stat::Min), Precision::exact(5));
924    }
925
926    #[test]
927    fn merge_min_bound_bound_lower() {
928        let merged = StatsSet::of(Stat::Min, Precision::inexact(4)).merge_ordered(
929            &StatsSet::of(Stat::Min, Precision::exact(5)),
930            &DType::Primitive(PType::I32, Nullability::NonNullable),
931        );
932        let merged_ref =
933            merged.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
934        assert_eq!(merged_ref.get_as::<i32>(Stat::Min), Precision::inexact(4));
935    }
936
937    #[test]
938    fn test_combine_is_constant() {
939        {
940            let mut stats = StatsSet::of(Stat::IsConstant, Precision::exact(true));
941            let stats2 = StatsSet::of(Stat::IsConstant, Precision::exact(true));
942            let mut stats_ref =
943                stats.as_mut_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
944            stats_ref
945                .combine_bool_stat::<IsConstant>(
946                    &stats2.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable)),
947                )
948                .unwrap();
949            assert_eq!(
950                stats_ref.get_as::<bool>(Stat::IsConstant),
951                Precision::exact(true)
952            );
953        }
954
955        {
956            let mut stats = StatsSet::of(Stat::IsConstant, Precision::exact(true));
957            let stats2 = StatsSet::of(Stat::IsConstant, Precision::inexact(false));
958            let mut stats_ref =
959                stats.as_mut_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
960            stats_ref
961                .combine_bool_stat::<IsConstant>(
962                    &stats2.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable)),
963                )
964                .unwrap();
965            assert_eq!(
966                stats_ref.get_as::<bool>(Stat::IsConstant),
967                Precision::exact(true)
968            );
969        }
970
971        {
972            let mut stats = StatsSet::of(Stat::IsConstant, Precision::exact(false));
973            let stats2 = StatsSet::of(Stat::IsConstant, Precision::inexact(false));
974            let mut stats_ref =
975                stats.as_mut_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
976            stats_ref
977                .combine_bool_stat::<IsConstant>(
978                    &stats2.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable)),
979                )
980                .unwrap();
981            assert_eq!(
982                stats_ref.get_as::<bool>(Stat::IsConstant),
983                Precision::exact(false)
984            );
985        }
986    }
987
988    #[test]
989    fn test_combine_sets_boolean_conflict() {
990        let mut stats1 = StatsSet::from_iter([
991            (Stat::IsConstant, Precision::exact(true)),
992            (Stat::IsSorted, Precision::exact(true)),
993        ]);
994
995        let stats2 = StatsSet::from_iter([
996            (Stat::IsConstant, Precision::exact(false)),
997            (Stat::IsSorted, Precision::exact(true)),
998        ]);
999
1000        let result = stats1.combine_sets(
1001            &stats2,
1002            &DType::Primitive(PType::I32, Nullability::NonNullable),
1003        );
1004        assert!(result.is_err());
1005    }
1006
1007    #[test]
1008    fn test_combine_sets_with_missing_stats() {
1009        let mut stats1 = StatsSet::from_iter([
1010            (Stat::Min, Precision::exact(42)),
1011            (Stat::UncompressedSizeInBytes, Precision::exact(1000)),
1012        ]);
1013
1014        let stats2 = StatsSet::from_iter([
1015            (Stat::Max, Precision::exact(100)),
1016            (Stat::IsStrictSorted, Precision::exact(true)),
1017        ]);
1018
1019        stats1
1020            .combine_sets(
1021                &stats2,
1022                &DType::Primitive(PType::I32, Nullability::NonNullable),
1023            )
1024            .unwrap();
1025
1026        let stats_ref =
1027            stats1.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
1028
1029        // Min should remain unchanged
1030        assert_eq!(stats_ref.get_as::<i32>(Stat::Min), Precision::exact(42));
1031        // Max should be added
1032        assert_eq!(stats_ref.get_as::<i32>(Stat::Max), Precision::exact(100));
1033        // IsStrictSorted should be added
1034        assert_eq!(
1035            stats_ref.get_as::<bool>(Stat::IsStrictSorted),
1036            Precision::exact(true)
1037        );
1038    }
1039
1040    #[test]
1041    fn test_combine_sets_with_inexact() {
1042        let mut stats1 = StatsSet::from_iter([
1043            (Stat::Min, Precision::exact(42)),
1044            (Stat::Max, Precision::inexact(100)),
1045            (Stat::IsConstant, Precision::exact(false)),
1046        ]);
1047
1048        let stats2 = StatsSet::from_iter([
1049            // Must ensure Min from stats2 is <= Min from stats1
1050            (Stat::Min, Precision::inexact(40)),
1051            (Stat::Max, Precision::exact(90)),
1052            (Stat::IsSorted, Precision::exact(true)),
1053        ]);
1054
1055        stats1
1056            .combine_sets(
1057                &stats2,
1058                &DType::Primitive(PType::I32, Nullability::NonNullable),
1059            )
1060            .unwrap();
1061
1062        let stats_ref =
1063            stats1.as_typed_ref(&DType::Primitive(PType::I32, Nullability::NonNullable));
1064
1065        // Min should remain unchanged since it's more restrictive than the inexact value
1066        assert_eq!(stats_ref.get_as::<i32>(Stat::Min), Precision::exact(42));
1067        // Check that max was updated with the exact value
1068        assert_eq!(stats_ref.get_as::<i32>(Stat::Max), Precision::exact(90));
1069        // Check that IsSorted was added
1070        assert_eq!(
1071            stats_ref.get_as::<bool>(Stat::IsSorted),
1072            Precision::exact(true)
1073        );
1074    }
1075}
vortex_array/stats/stats_set.rs

vortex_array/stats/
stats_set.rs
