vortex_array/stats/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Traits and utilities to compute and access array statistics.
5
6use arrow_buffer::BooleanBufferBuilder;
7use arrow_buffer::MutableBuffer;
8use arrow_buffer::bit_iterator::BitIterator;
9use enum_iterator::last;
10pub use expr::all_nan;
11pub use expr::all_non_nan;
12pub use expr::all_non_null;
13pub use expr::all_null;
14pub use expr::bound;
15pub use expr::min_max;
16pub use expr::nan_count;
17pub use expr::null_count;
18pub use expr::stat;
19pub use expr::sum;
20pub use stats_set::*;
21
22mod array;
23pub mod bind;
24pub mod expr;
25pub mod flatbuffers;
26pub mod rewrite;
27pub mod session;
28mod stats_set;
29
30pub use array::*;
31pub use session::*;
32use vortex_error::VortexExpect;
33
34use crate::expr::stats::Stat;
35
36/// Statistics that are used for pruning files (i.e., we want to ensure they are computed when compressing/writing).
37/// Sum is included for boolean arrays.
38pub const PRUNING_STATS: &[Stat] = &[
39    Stat::Min,
40    Stat::Max,
41    Stat::Sum,
42    Stat::NullCount,
43    Stat::NaNCount,
44];
45
46pub fn as_stat_bitset_bytes(stats: &[Stat]) -> Vec<u8> {
47    let max_stat = u8::from(last::<Stat>().vortex_expect("last stat")) as usize + 1;
48    // TODO(ngates): use vortex-buffer::BitBuffer
49    let mut stat_bitset = BooleanBufferBuilder::new_from_buffer(
50        MutableBuffer::from_len_zeroed(max_stat.div_ceil(8)),
51        max_stat,
52    );
53    for stat in stats {
54        stat_bitset.set_bit(u8::from(*stat) as usize, true);
55    }
56
57    stat_bitset
58        .finish()
59        .into_inner()
60        .into_vec()
61        .unwrap_or_else(|b| b.to_vec())
62}
63
64pub fn stats_from_bitset_bytes(bytes: &[u8]) -> Vec<Stat> {
65    BitIterator::new(bytes, 0, bytes.len() * 8)
66        .enumerate()
67        .filter_map(|(i, b)| b.then_some(i))
68        // Filter out indices failing conversion, these are stats written by newer version of library
69        .filter_map(|i| {
70            let Ok(stat) = u8::try_from(i) else {
71                tracing::debug!("invalid stat encountered: {i}");
72                return None;
73            };
74            Stat::try_from(stat).ok()
75        })
76        .collect::<Vec<_>>()
77}
vortex_array/stats/mod.rs

vortex_array/stats/
mod.rs
