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

vortex_array/stats/
array.rs
