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;
7
8use parking_lot::RwLock;
9use vortex_error::{VortexError, VortexResult, vortex_panic};
10use vortex_scalar::ScalarValue;
11
12use super::{
13    Precision, Stat, StatType, StatsProvider, StatsProviderExt, StatsSet, StatsSetIntoIter,
14};
15use crate::Array;
16use crate::compute::{
17    MinMaxResult, is_constant, is_sorted, is_strict_sorted, min_max, nan_count, sum,
18};
19
20/// A shared [`StatsSet`] stored in an array. Can be shared by copies of the array and can also be mutated in place.
21// TODO(adamg): This is a very bad name.
22#[derive(Clone, Default, Debug)]
23pub struct ArrayStats {
24    inner: Arc<RwLock<StatsSet>>,
25}
26
27/// Reference to an array's [`StatsSet`]. Can be used to get and mutate the underlying stats.
28///
29/// Constructed by calling [`ArrayStats::to_ref`].
30pub struct StatsSetRef<'a> {
31    // We need to reference back to the array
32    dyn_array_ref: &'a dyn Array,
33    parent_stats: ArrayStats,
34}
35
36impl ArrayStats {
37    pub fn to_ref<'a>(&self, array: &'a dyn Array) -> StatsSetRef<'a> {
38        StatsSetRef {
39            dyn_array_ref: array,
40            parent_stats: self.clone(),
41        }
42    }
43
44    pub fn set(&self, stat: Stat, value: Precision<ScalarValue>) {
45        self.inner.write().set(stat, value);
46    }
47
48    pub fn clear(&self, stat: Stat) {
49        self.inner.write().clear(stat);
50    }
51
52    pub fn retain(&self, stats: &[Stat]) {
53        self.inner.write().retain_only(stats);
54    }
55}
56
57impl From<StatsSet> for ArrayStats {
58    fn from(value: StatsSet) -> Self {
59        Self {
60            inner: Arc::new(RwLock::new(value)),
61        }
62    }
63}
64
65impl From<ArrayStats> for StatsSet {
66    fn from(value: ArrayStats) -> Self {
67        value.inner.read().clone()
68    }
69}
70
71impl StatsProvider for ArrayStats {
72    fn get(&self, stat: Stat) -> Option<Precision<ScalarValue>> {
73        let guard = self.inner.read();
74        guard.get(stat)
75    }
76
77    fn len(&self) -> usize {
78        let guard = self.inner.read();
79        guard.len()
80    }
81}
82
83impl StatsSetRef<'_> {
84    pub fn set_iter(&self, iter: StatsSetIntoIter) {
85        let mut guard = self.parent_stats.inner.write();
86
87        for (stat, value) in iter {
88            guard.set(stat, value);
89        }
90    }
91
92    pub fn inherit(&self, parent_stats: StatsSetRef<'_>) {
93        // TODO(ngates): depending on statistic, this should choose the more precise one
94        self.set_iter(parent_stats.into_iter());
95    }
96
97    // TODO(adamg): potentially problematic name
98    pub fn to_owned(&self) -> StatsSet {
99        self.parent_stats.inner.read().clone()
100    }
101
102    pub fn into_iter(&self) -> StatsSetIntoIter {
103        self.to_owned().into_iter()
104    }
105
106    pub fn compute_stat(&self, stat: Stat) -> VortexResult<Option<ScalarValue>> {
107        // If it's already computed and exact, we can return it.
108        if let Some(Precision::Exact(stat)) = self.get(stat) {
109            return Ok(Some(stat));
110        }
111
112        Ok(match stat {
113            Stat::Min => {
114                min_max(self.dyn_array_ref)?.map(|MinMaxResult { min, max: _ }| min.into_value())
115            }
116            Stat::Max => {
117                min_max(self.dyn_array_ref)?.map(|MinMaxResult { min: _, max }| max.into_value())
118            }
119            Stat::Sum => {
120                Stat::Sum
121                    .dtype(self.dyn_array_ref.dtype())
122                    .is_some()
123                    .then(|| {
124                        // Sum is supported for this dtype.
125                        sum(self.dyn_array_ref)
126                    })
127                    .transpose()?
128                    .map(|s| s.into_value())
129            }
130            Stat::NullCount => Some(self.dyn_array_ref.invalid_count()?.into()),
131            Stat::IsConstant => {
132                if self.dyn_array_ref.is_empty() {
133                    None
134                } else {
135                    is_constant(self.dyn_array_ref)?.map(ScalarValue::from)
136                }
137            }
138            Stat::IsSorted => Some(is_sorted(self.dyn_array_ref)?.into()),
139            Stat::IsStrictSorted => Some(is_strict_sorted(self.dyn_array_ref)?.into()),
140            Stat::UncompressedSizeInBytes => {
141                let nbytes: ScalarValue =
142                    self.dyn_array_ref.to_canonical()?.as_ref().nbytes().into();
143                self.set(stat, Precision::exact(nbytes.clone()));
144                Some(nbytes)
145            }
146            Stat::NaNCount => {
147                Stat::NaNCount
148                    .dtype(self.dyn_array_ref.dtype())
149                    .is_some()
150                    .then(|| {
151                        // NaNCount is supported for this dtype.
152                        nan_count(self.dyn_array_ref)
153                    })
154                    .transpose()?
155                    .map(|s| s.into())
156            }
157        })
158    }
159
160    pub fn compute_all(&self, stats: &[Stat]) -> VortexResult<StatsSet> {
161        let mut stats_set = StatsSet::default();
162        for &stat in stats {
163            if let Some(s) = self.compute_stat(stat)? {
164                stats_set.set(stat, Precision::exact(s))
165            }
166        }
167        Ok(stats_set)
168    }
169}
170
171impl StatsSetRef<'_> {
172    pub fn get_as<U: for<'a> TryFrom<&'a ScalarValue, Error = VortexError>>(
173        &self,
174        stat: Stat,
175    ) -> Option<Precision<U>> {
176        StatsProviderExt::get_as::<U>(self, stat)
177    }
178
179    pub fn get_as_bound<S, U>(&self) -> Option<S::Bound>
180    where
181        S: StatType<U>,
182        U: for<'a> TryFrom<&'a ScalarValue, Error = VortexError>,
183    {
184        StatsProviderExt::get_as_bound::<S, U>(self)
185    }
186
187    pub fn compute_as<U: for<'a> TryFrom<&'a ScalarValue, Error = VortexError>>(
188        &self,
189        stat: Stat,
190    ) -> Option<U> {
191        self.compute_stat(stat)
192            .inspect_err(|e| log::warn!("Failed to compute stat {stat}: {e}"))
193            .ok()
194            .flatten()
195            .map(|s| U::try_from(&s))
196            .transpose()
197            .unwrap_or_else(|err| {
198                vortex_panic!(
199                    err,
200                    "Failed to compute stat {} as {}",
201                    stat,
202                    std::any::type_name::<U>()
203                )
204            })
205    }
206
207    pub fn set(&self, stat: Stat, value: Precision<ScalarValue>) {
208        self.parent_stats.set(stat, value);
209    }
210
211    pub fn clear(&self, stat: Stat) {
212        self.parent_stats.clear(stat);
213    }
214
215    pub fn retain(&self, stats: &[Stat]) {
216        self.parent_stats.retain(stats);
217    }
218
219    pub fn compute_min<U: for<'a> TryFrom<&'a ScalarValue, Error = VortexError>>(
220        &self,
221    ) -> Option<U> {
222        self.compute_as(Stat::Min)
223    }
224
225    pub fn compute_max<U: for<'a> TryFrom<&'a ScalarValue, Error = VortexError>>(
226        &self,
227    ) -> Option<U> {
228        self.compute_as(Stat::Max)
229    }
230
231    pub fn compute_is_sorted(&self) -> Option<bool> {
232        self.compute_as(Stat::IsSorted)
233    }
234
235    pub fn compute_is_strict_sorted(&self) -> Option<bool> {
236        self.compute_as(Stat::IsStrictSorted)
237    }
238
239    pub fn compute_is_constant(&self) -> Option<bool> {
240        self.compute_as(Stat::IsConstant)
241    }
242
243    pub fn compute_null_count(&self) -> Option<usize> {
244        self.compute_as(Stat::NullCount)
245    }
246
247    pub fn compute_uncompressed_size_in_bytes(&self) -> Option<usize> {
248        self.compute_as(Stat::UncompressedSizeInBytes)
249    }
250}
251
252impl StatsProvider for StatsSetRef<'_> {
253    fn get(&self, stat: Stat) -> Option<Precision<ScalarValue>> {
254        self.parent_stats.get(stat)
255    }
256
257    fn len(&self) -> usize {
258        self.parent_stats.len()
259    }
260}
vortex_array/stats/array.rs

vortex_array/stats/
array.rs
