vortex_array/stats/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Stats as they are stored on arrays.
5
6use std::sync::Arc;
7use std::sync::OnceLock;
8
9use parking_lot::RwLock;
10use vortex_array::ExecutionCtx;
11use vortex_error::VortexError;
12use vortex_error::VortexResult;
13use vortex_error::vortex_panic;
14
15use super::MutTypedStatsSetRef;
16use super::StatsSet;
17use super::StatsSetIntoIter;
18use super::TypedStatsSetRef;
19use crate::ArrayRef;
20use crate::aggregate_fn::NumericalAggregateOpts;
21use crate::aggregate_fn::fns::is_constant::is_constant;
22use crate::aggregate_fn::fns::is_sorted::is_sorted;
23use crate::aggregate_fn::fns::is_sorted::is_strict_sorted;
24use crate::aggregate_fn::fns::min_max::MinMaxResult;
25use crate::aggregate_fn::fns::min_max::min_max;
26use crate::aggregate_fn::fns::nan_count::nan_count;
27use crate::aggregate_fn::fns::sum::sum;
28use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeOpts;
29use crate::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes;
30use crate::expr::stats::Precision;
31use crate::expr::stats::Stat;
32use crate::expr::stats::StatsProvider;
33use crate::scalar::Scalar;
34use crate::scalar::ScalarValue;
35
36/// A shared [`StatsSet`] stored in an array. Can be shared by copies of the array and can also be mutated in place.
37/// Stats are created lazily on read or write.
38// TODO(adamg): This is a very bad name.
39#[derive(Default, Debug)]
40pub struct ArrayStats {
41    inner: OnceLock<Arc<RwLock<StatsSet>>>,
42}
43
44impl Clone for ArrayStats {
45    fn clone(&self) -> Self {
46        Self {
47            inner: OnceLock::from(Arc::clone(self.shared())),
48        }
49    }
50}
51
52/// Reference to an array's [`StatsSet`]. Can be used to get and mutate the underlying stats.
53///
54/// Constructed by calling [`ArrayStats::to_ref`].
55pub struct StatsSetRef<'a> {
56    // We need to reference back to the array
57    dyn_array_ref: &'a ArrayRef,
58    array_stats: &'a ArrayStats,
59}
60
61impl ArrayStats {
62    pub fn to_ref<'a>(&'a self, array: &'a ArrayRef) -> StatsSetRef<'a> {
63        StatsSetRef {
64            dyn_array_ref: array,
65            array_stats: self,
66        }
67    }
68
69    fn shared(&self) -> &Arc<RwLock<StatsSet>> {
70        self.inner.get_or_init(Default::default)
71    }
72
73    fn read<R>(&self, f: impl FnOnce(&StatsSet) -> R) -> R {
74        match self.inner.get() {
75            Some(shared) => f(&shared.read()),
76            None => f(&StatsSet::default()),
77        }
78    }
79
80    pub fn set(&self, stat: Stat, value: Precision<ScalarValue>) {
81        self.shared().write().set(stat, value);
82    }
83
84    pub fn clear(&self, stat: Stat) {
85        if let Some(shared) = self.inner.get() {
86            shared.write().clear(stat);
87        }
88    }
89
90    pub fn retain(&self, stats: &[Stat]) {
91        if let Some(shared) = self.inner.get() {
92            shared.write().retain_only(stats);
93        }
94    }
95}
96
97impl From<StatsSet> for ArrayStats {
98    fn from(value: StatsSet) -> Self {
99        Self {
100            inner: OnceLock::from(Arc::new(RwLock::new(value))),
101        }
102    }
103}
104
105impl From<ArrayStats> for StatsSet {
106    fn from(value: ArrayStats) -> Self {
107        value.read(|stats| stats.clone())
108    }
109}
110
111impl StatsSetRef<'_> {
112    pub(crate) fn replace(&self, stats: StatsSet) {
113        if stats.is_empty() && self.array_stats.inner.get().is_none() {
114            return;
115        }
116        *self.array_stats.shared().write() = stats;
117    }
118
119    /// Set every stat in `stats`, overwriting existing values for the same stats.
120    ///
121    /// When the array has no stats yet, the set is moved in whole rather than merged one stat at
122    /// a time.
123    pub fn set_all(&self, stats: StatsSet) {
124        if stats.is_empty() {
125            return;
126        }
127        let mut guard = self.array_stats.shared().write();
128        if guard.is_empty() {
129            *guard = stats;
130        } else {
131            for (stat, value) in stats {
132                guard.set(stat, value);
133            }
134        }
135    }
136
137    pub fn set_iter(&self, iter: StatsSetIntoIter) {
138        let mut iter = iter.peekable();
139        if iter.peek().is_none() {
140            return;
141        }
142        let mut guard = self.array_stats.shared().write();
143        for (stat, value) in iter {
144            guard.set(stat, value);
145        }
146    }
147
148    pub fn inherit_from(&self, stats: StatsSetRef<'_>) {
149        let Some(source) = stats.array_stats.inner.get() else {
150            return;
151        };
152        // Only inherit if the underlying stats are different
153        if self
154            .array_stats
155            .inner
156            .get()
157            .is_none_or(|shared| !Arc::ptr_eq(shared, source))
158        {
159            stats.with_iter(|iter| self.inherit(iter));
160        }
161    }
162
163    pub fn inherit<'a>(&self, iter: impl Iterator<Item = &'a (Stat, Precision<ScalarValue>)>) {
164        let mut iter = iter.peekable();
165        if iter.peek().is_none() {
166            return;
167        }
168        let mut guard = self.array_stats.shared().write();
169        for (stat, value) in iter {
170            if !value.is_exact() {
171                if !guard.get(*stat).is_exact() {
172                    guard.set(*stat, value.clone());
173                }
174            } else {
175                guard.set(*stat, value.clone());
176            }
177        }
178    }
179
180    pub fn with_typed_stats_set<U, F: FnOnce(TypedStatsSetRef) -> U>(&self, apply: F) -> U {
181        self.array_stats
182            .read(|stats| apply(stats.as_typed_ref(self.dyn_array_ref.dtype())))
183    }
184
185    pub fn with_mut_typed_stats_set<U, F: FnOnce(MutTypedStatsSetRef) -> U>(&self, apply: F) -> U {
186        apply(
187            self.array_stats
188                .shared()
189                .write()
190                .as_mut_typed_ref(self.dyn_array_ref.dtype()),
191        )
192    }
193
194    pub fn to_owned(&self) -> StatsSet {
195        self.array_stats.read(|stats| stats.clone())
196    }
197
198    /// Returns a clone of the underlying [`ArrayStats`].
199    pub fn to_array_stats(&self) -> ArrayStats {
200        self.array_stats.clone()
201    }
202
203    /// Share underlying stats if they exist or return a lazy instance
204    pub(crate) fn share_existing(&self) -> ArrayStats {
205        ArrayStats {
206            inner: match self.array_stats.inner.get() {
207                Some(shared) => OnceLock::from(Arc::clone(shared)),
208                None => OnceLock::new(),
209            },
210        }
211    }
212
213    pub fn with_iter<
214        F: for<'a> FnOnce(&mut dyn Iterator<Item = &'a (Stat, Precision<ScalarValue>)>) -> R,
215        R,
216    >(
217        &self,
218        f: F,
219    ) -> R {
220        self.array_stats.read(|stats| f(&mut stats.iter()))
221    }
222
223    /// Returns the value of `stat` by either fetching it from cache if it exists and is [`Precision::Exact`], or falling back to
224    /// computation. The underlying compute kernels will cache the computed stat in the latter case.
225    pub fn compute_stat(&self, stat: Stat, ctx: &mut ExecutionCtx) -> VortexResult<Option<Scalar>> {
226        // If it's already computed and exact, we can return it.
227        if let Precision::Exact(s) = self.get(stat) {
228            return Ok(Some(s));
229        }
230
231        Ok(match stat {
232            Stat::Min => min_max(self.dyn_array_ref, ctx, NumericalAggregateOpts::default())?
233                .map(|MinMaxResult { min, max: _ }| min),
234            Stat::Max => min_max(self.dyn_array_ref, ctx, NumericalAggregateOpts::default())?
235                .map(|MinMaxResult { min: _, max }| max),
236            Stat::Sum => {
237                Stat::Sum
238                    .dtype(self.dyn_array_ref.dtype())
239                    .is_some()
240                    .then(|| {
241                        // Sum is supported for this dtype.
242                        sum(self.dyn_array_ref, ctx)
243                    })
244                    .transpose()?
245            }
246            Stat::NullCount => self.dyn_array_ref.invalid_count(ctx).ok().map(Into::into),
247            Stat::IsConstant => {
248                if self.dyn_array_ref.is_empty() {
249                    None
250                } else {
251                    Some(is_constant(self.dyn_array_ref, ctx)?.into())
252                }
253            }
254            Stat::IsSorted => Some(is_sorted(self.dyn_array_ref, ctx)?.into()),
255            Stat::IsStrictSorted => Some(is_strict_sorted(self.dyn_array_ref, ctx)?.into()),
256            Stat::UncompressedSizeInBytes => Stat::UncompressedSizeInBytes
257                .dtype(self.dyn_array_ref.dtype())
258                .is_some()
259                .then(|| {
260                    uncompressed_size_in_bytes(
261                        self.dyn_array_ref,
262                        UncompressedSizeOpts::exact(),
263                        ctx,
264                    )
265                })
266                .transpose()?
267                .map(|s| s.into()),
268            Stat::NaNCount => {
269                Stat::NaNCount
270                    .dtype(self.dyn_array_ref.dtype())
271                    .is_some()
272                    .then(|| {
273                        // NaNCount is supported for this dtype.
274                        nan_count(self.dyn_array_ref, ctx)
275                    })
276                    .transpose()?
277                    .map(|s| s.into())
278            }
279        })
280    }
281
282    pub fn compute_all(&self, stats: &[Stat], ctx: &mut ExecutionCtx) -> VortexResult<StatsSet> {
283        let mut stats_set = StatsSet::default();
284        for &stat in stats {
285            if let Some(s) = self.compute_stat(stat, ctx)?
286                && let Some(value) = s.into_value()
287            {
288                stats_set.set(stat, Precision::exact(value));
289            }
290        }
291        Ok(stats_set)
292    }
293}
294
295impl StatsSetRef<'_> {
296    pub fn compute_as<U: for<'a> TryFrom<&'a Scalar, Error = VortexError>>(
297        &self,
298        stat: Stat,
299        ctx: &mut ExecutionCtx,
300    ) -> Option<U> {
301        self.compute_stat(stat, ctx)
302            .inspect_err(|e| tracing::warn!("Failed to compute stat {stat}: {e}"))
303            .ok()
304            .flatten()
305            .map(|s| U::try_from(&s))
306            .transpose()
307            .unwrap_or_else(|err| {
308                vortex_panic!(
309                    err,
310                    "Failed to compute stat {} as {}",
311                    stat,
312                    std::any::type_name::<U>()
313                )
314            })
315    }
316
317    pub fn set(&self, stat: Stat, value: Precision<ScalarValue>) {
318        self.array_stats.set(stat, value);
319    }
320
321    pub fn clear(&self, stat: Stat) {
322        self.array_stats.clear(stat);
323    }
324
325    pub fn compute_min<U: for<'a> TryFrom<&'a Scalar, Error = VortexError>>(
326        &self,
327        ctx: &mut ExecutionCtx,
328    ) -> Option<U> {
329        self.compute_as(Stat::Min, ctx)
330    }
331
332    pub fn compute_max<U: for<'a> TryFrom<&'a Scalar, Error = VortexError>>(
333        &self,
334        ctx: &mut ExecutionCtx,
335    ) -> Option<U> {
336        self.compute_as(Stat::Max, ctx)
337    }
338
339    pub fn compute_is_sorted(&self, ctx: &mut ExecutionCtx) -> Option<bool> {
340        self.compute_bool(Stat::IsSorted, ctx)
341    }
342
343    pub fn compute_is_strict_sorted(&self, ctx: &mut ExecutionCtx) -> Option<bool> {
344        self.compute_bool(Stat::IsStrictSorted, ctx)
345    }
346
347    pub fn compute_is_constant(&self, ctx: &mut ExecutionCtx) -> Option<bool> {
348        self.compute_bool(Stat::IsConstant, ctx)
349    }
350
351    pub fn compute_null_count(&self, ctx: &mut ExecutionCtx) -> Option<usize> {
352        self.compute_usize(Stat::NullCount, ctx)
353    }
354
355    pub fn compute_uncompressed_size_in_bytes(&self, ctx: &mut ExecutionCtx) -> Option<usize> {
356        self.compute_usize(Stat::UncompressedSizeInBytes, ctx)
357    }
358
359    /// Like [`Self::compute_as`], but reads a cached exact value without building a [`Scalar`].
360    fn compute_bool(&self, stat: Stat, ctx: &mut ExecutionCtx) -> Option<bool> {
361        // Bind the cached value first, so the read lock is released before any computation.
362        let cached = self.array_stats.read(|stats| stats.get_bool(stat));
363        if let Precision::Exact(value) = cached {
364            return Some(value);
365        }
366
367        self.compute_as(stat, ctx)
368    }
369
370    /// Like [`Self::compute_as`], but reads a cached exact value without building a [`Scalar`].
371    fn compute_usize(&self, stat: Stat, ctx: &mut ExecutionCtx) -> Option<usize> {
372        // Bind the cached value first, so the read lock is released before any computation.
373        let cached = self.array_stats.read(|stats| stats.get_usize(stat));
374        if let Precision::Exact(value) = cached {
375            return Some(value);
376        }
377
378        self.compute_as(stat, ctx)
379    }
380}
381
382impl StatsProvider for StatsSetRef<'_> {
383    fn get(&self, stat: Stat) -> Precision<Scalar> {
384        self.array_stats
385            .read(|stats| stats.as_typed_ref(self.dyn_array_ref.dtype()).get(stat))
386    }
387
388    fn len(&self) -> usize {
389        self.array_stats.read(|stats| stats.len())
390    }
391}
392
393#[cfg(test)]
394mod tests {
395    use super::ArrayStats;
396    use super::StatsSet;
397    use crate::expr::stats::Precision;
398    use crate::expr::stats::Stat;
399    use crate::scalar::ScalarValue;
400
401    #[test]
402    fn empty_stats() {
403        let stats = ArrayStats::default();
404        assert_eq!(StatsSet::from(stats).len(), 0);
405    }
406
407    #[test]
408    fn clone_stats() {
409        let stats = ArrayStats::default();
410        let copy = stats.clone();
411
412        stats.set(Stat::NullCount, Precision::exact(ScalarValue::from(3u64)));
413        let through_copy = StatsSet::from(copy);
414        assert_eq!(
415            through_copy.get(Stat::NullCount),
416            Precision::exact(ScalarValue::from(3u64))
417        );
418    }
419
420    #[test]
421    fn clone_writes() {
422        let stats = ArrayStats::default();
423        let copy = stats.clone();
424
425        copy.set(Stat::NullCount, Precision::exact(ScalarValue::from(7u64)));
426        assert_eq!(
427            StatsSet::from(stats).get(Stat::NullCount),
428            Precision::exact(ScalarValue::from(7u64))
429        );
430    }
431}
vortex_array/stats/array.rs

vortex_array/stats/
array.rs
