vortex_array/expr/stats/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Debug;
5use std::fmt::Display;
6use std::fmt::Formatter;
7
8use enum_iterator::Sequence;
9use enum_iterator::all;
10use num_enum::IntoPrimitive;
11use num_enum::TryFromPrimitive;
12
13use crate::dtype::DType;
14use crate::dtype::Nullability::NonNullable;
15
16mod bound;
17mod precision;
18mod provider;
19mod stat_bound;
20
21pub use bound::*;
22pub use precision::*;
23pub use provider::*;
24pub use stat_bound::*;
25
26use crate::aggregate_fn;
27use crate::aggregate_fn::AggregateFnRef;
28use crate::aggregate_fn::AggregateFnVTable;
29use crate::aggregate_fn::AggregateFnVTableExt;
30use crate::aggregate_fn::EmptyOptions;
31use crate::aggregate_fn::NumericalAggregateOpts;
32use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeOpts;
33
34#[derive(
35    Debug,
36    Clone,
37    Copy,
38    PartialEq,
39    Eq,
40    PartialOrd,
41    Ord,
42    Hash,
43    Sequence,
44    IntoPrimitive,
45    TryFromPrimitive,
46)]
47#[repr(u8)]
48pub enum Stat {
49    /// Whether all values are the same (nulls are not equal to other non-null values,
50    /// so this is true iff all values are null or all values are the same non-null value)
51    IsConstant = 0,
52    /// Whether the non-null values in the array are sorted in ascending order (i.e., we skip nulls)
53    /// This may later be extended to support descending order, but for now we only support ascending order.
54    IsSorted = 1,
55    /// Whether the non-null values in the array are strictly sorted in ascending order (i.e., sorted with no duplicates)
56    /// This may later be extended to support descending order, but for now we only support ascending order.
57    IsStrictSorted = 2,
58    /// The maximum value in the array (ignoring nulls, unless all values are null)
59    Max = 3,
60    /// The minimum value in the array (ignoring nulls, unless all values are null)
61    Min = 4,
62    /// The sum of the non-null values of the array.
63    Sum = 5,
64    /// The number of null values in the array
65    NullCount = 6,
66    /// The uncompressed size of the array in bytes
67    UncompressedSizeInBytes = 7,
68    /// The number of NaN values in the array
69    NaNCount = 8,
70}
71
72/// These structs allow the extraction of the bound from the `Precision` value.
73/// They tie together the Stat and the StatBound, which allows the bound to be extracted.
74pub struct Max;
75
76pub struct Min;
77
78pub struct Sum;
79
80pub struct IsConstant;
81
82pub struct IsSorted;
83
84pub struct IsStrictSorted;
85
86pub struct NullCount;
87
88pub struct UncompressedSizeInBytes;
89
90pub struct NaNCount;
91
92impl StatType<bool> for IsConstant {
93    type Bound = Precision<bool>;
94
95    const STAT: Stat = Stat::IsConstant;
96}
97
98impl StatType<bool> for IsSorted {
99    type Bound = Precision<bool>;
100
101    const STAT: Stat = Stat::IsSorted;
102}
103
104impl StatType<bool> for IsStrictSorted {
105    type Bound = Precision<bool>;
106
107    const STAT: Stat = Stat::IsStrictSorted;
108}
109
110impl<T: PartialOrd + Clone> StatType<T> for NullCount {
111    type Bound = UpperBound<T>;
112
113    const STAT: Stat = Stat::NullCount;
114}
115
116impl<T: PartialOrd + Clone> StatType<T> for UncompressedSizeInBytes {
117    type Bound = UpperBound<T>;
118
119    const STAT: Stat = Stat::UncompressedSizeInBytes;
120}
121
122impl<T: PartialOrd + Clone + Debug> StatType<T> for Max {
123    type Bound = UpperBound<T>;
124
125    const STAT: Stat = Stat::Max;
126}
127
128impl<T: PartialOrd + Clone + Debug> StatType<T> for Min {
129    type Bound = LowerBound<T>;
130
131    const STAT: Stat = Stat::Min;
132}
133
134impl<T: PartialOrd + Clone + Debug> StatType<T> for Sum {
135    type Bound = Precision<T>;
136
137    const STAT: Stat = Stat::Sum;
138}
139
140impl<T: PartialOrd + Clone> StatType<T> for NaNCount {
141    type Bound = UpperBound<T>;
142
143    const STAT: Stat = Stat::NaNCount;
144}
145
146impl Stat {
147    /// Whether the statistic is commutative (i.e., whether merging can be done independently of ordering)
148    /// e.g., min/max are commutative, but is_sorted is not
149    pub fn is_commutative(&self) -> bool {
150        // NOTE: we prefer this syntax to force a compile error if we add a new stat
151        match self {
152            Self::IsConstant
153            | Self::Max
154            | Self::Min
155            | Self::NullCount
156            | Self::Sum
157            | Self::NaNCount
158            | Self::UncompressedSizeInBytes => true,
159            Self::IsSorted | Self::IsStrictSorted => false,
160        }
161    }
162
163    /// Whether the statistic has the same dtype as the array it's computed on
164    pub fn has_same_dtype_as_array(&self) -> bool {
165        matches!(self, Stat::Min | Stat::Max)
166    }
167
168    /// Return the [`DType`] of the statistic scalar assuming the array is of the given [`DType`].
169    pub fn dtype(&self, data_type: &DType) -> Option<DType> {
170        Some(match self {
171            Self::IsConstant => DType::Bool(NonNullable),
172            Self::IsSorted => DType::Bool(NonNullable),
173            Self::IsStrictSorted => DType::Bool(NonNullable),
174            Self::Max if matches!(data_type, DType::Null) => return None,
175            Self::Max => data_type.clone(),
176            Self::Min if matches!(data_type, DType::Null) => return None,
177            Self::Min => data_type.clone(),
178            Self::NullCount => {
179                return aggregate_fn::fns::null_count::NullCount
180                    .return_dtype(&EmptyOptions, data_type);
181            }
182            Self::UncompressedSizeInBytes => {
183                return aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes
184                    .return_dtype(&UncompressedSizeOpts::exact(), data_type);
185            }
186            Self::NaNCount => {
187                return aggregate_fn::fns::nan_count::NanCount
188                    .return_dtype(&EmptyOptions, data_type);
189            }
190            Self::Sum => {
191                // Statistics follow NaN-skipping semantics; request it explicitly.
192                return aggregate_fn::fns::sum::Sum
193                    .return_dtype(&NumericalAggregateOpts::skip_nans(), data_type);
194            }
195        })
196    }
197
198    /// Return the built-in aggregate function corresponding to this statistic, if one exists.
199    pub fn aggregate_fn(&self) -> Option<AggregateFnRef> {
200        // Statistics follow NaN-skipping semantics; request it explicitly rather than the default.
201        Some(match self {
202            Self::Max => aggregate_fn::fns::max::Max.bind(NumericalAggregateOpts::skip_nans()),
203            Self::Min => aggregate_fn::fns::min::Min.bind(NumericalAggregateOpts::skip_nans()),
204            Self::Sum => aggregate_fn::fns::sum::Sum.bind(NumericalAggregateOpts::skip_nans()),
205            Self::NullCount => aggregate_fn::fns::null_count::NullCount.bind(EmptyOptions),
206            Self::NaNCount => aggregate_fn::fns::nan_count::NanCount.bind(EmptyOptions),
207            // The stat slot holds the exact size, so the statistic measures referenced bytes.
208            Self::UncompressedSizeInBytes => {
209                aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes
210                    .bind(UncompressedSizeOpts::exact())
211            }
212            Self::IsConstant | Self::IsSorted | Self::IsStrictSorted => return None,
213        })
214    }
215
216    /// Return the statistic represented by `aggregate_fn`, if it has a legacy stat slot.
217    ///
218    /// Min/max/sum statistics skip NaN values, so NaN-including configurations of those
219    /// aggregates have no stat slot. The uncompressed size statistic is exact, so the inexact
220    /// configuration of that aggregate has no stat slot either.
221    pub fn from_aggregate_fn(aggregate_fn: &AggregateFnRef) -> Option<Self> {
222        if let Some(options) = aggregate_fn.as_opt::<aggregate_fn::fns::sum::Sum>() {
223            return options.skip_nans.then_some(Self::Sum);
224        }
225        if aggregate_fn.is::<aggregate_fn::fns::nan_count::NanCount>() {
226            return Some(Self::NaNCount);
227        }
228        if aggregate_fn.is::<aggregate_fn::fns::null_count::NullCount>() {
229            return Some(Self::NullCount);
230        }
231        if let Some(options) = aggregate_fn.as_opt::<aggregate_fn::fns::min::Min>() {
232            return options.skip_nans.then_some(Self::Min);
233        }
234        if let Some(options) = aggregate_fn.as_opt::<aggregate_fn::fns::max::Max>() {
235            return options.skip_nans.then_some(Self::Max);
236        }
237        if let Some(options) =
238            aggregate_fn
239                .as_opt::<aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes>()
240        {
241            return options.is_exact().then_some(Self::UncompressedSizeInBytes);
242        }
243        None
244    }
245
246    pub fn name(&self) -> &str {
247        match self {
248            Self::IsConstant => "is_constant",
249            Self::IsSorted => "is_sorted",
250            Self::IsStrictSorted => "is_strict_sorted",
251            Self::Max => "max",
252            Self::Min => "min",
253            Self::NullCount => "null_count",
254            Self::UncompressedSizeInBytes => "uncompressed_size_in_bytes",
255            Self::Sum => "sum",
256            Self::NaNCount => "nan_count",
257        }
258    }
259
260    pub fn all() -> impl Iterator<Item = Stat> {
261        all::<Self>()
262    }
263}
264
265impl Display for Stat {
266    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
267        write!(f, "{}", self.name())
268    }
269}
270
271#[cfg(test)]
272mod test {
273    use enum_iterator::all;
274
275    use crate::VortexSessionExecute;
276    use crate::array_session;
277    use crate::arrays::PrimitiveArray;
278    use crate::expr::stats::Stat;
279
280    #[test]
281    fn min_of_nulls_is_not_panic() {
282        let min = PrimitiveArray::from_option_iter::<i32, _>([None, None, None, None])
283            .statistics()
284            .compute_as::<i64>(Stat::Min, &mut array_session().create_execution_ctx());
285
286        assert_eq!(min, None);
287    }
288
289    #[test]
290    fn has_same_dtype_as_array() {
291        assert!(Stat::Min.has_same_dtype_as_array());
292        assert!(Stat::Max.has_same_dtype_as_array());
293        for stat in all::<Stat>().filter(|s| !matches!(s, Stat::Min | Stat::Max)) {
294            assert!(!stat.has_same_dtype_as_array());
295        }
296    }
297}
vortex_array/expr/stats/mod.rs

vortex_array/expr/stats/
mod.rs
