Skip to main content

vortex_compressor/stats/
varbinview.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Variable-length byte/string compression statistics.
5
6use vortex_array::ExecutionCtx;
7use vortex_array::arrays::VarBinViewArray;
8use vortex_error::VortexExpect;
9use vortex_error::VortexResult;
10use vortex_error::vortex_err;
11use vortex_utils::aliases::hash_set::HashSet;
12
13use super::GenerateStatsOptions;
14
15/// Array of variable-length byte/string values, and relevant stats for compression.
16#[derive(Clone, Debug)]
17pub struct StringStats {
18    /// The estimated number of distinct values, or `None` if not computed.
19    /// This _must_ be non-zero.
20    estimated_distinct_count: Option<u32>,
21    /// The number of non-null values.
22    value_count: u32,
23    /// The number of null values.
24    null_count: u32,
25}
26
27/// Estimate the number of distinct values in the var bin view array.
28fn estimate_distinct_count(varbinview: &VarBinViewArray) -> VortexResult<u32> {
29    let views = varbinview.views();
30    // Equal values have equal heads, so distinct heads under-count distinct values.
31    // NOTE: there are cases where this performs pessimally, e.g. when we have strings that all
32    // share a 4-byte prefix and have the same length.
33    let mut distinct = HashSet::with_capacity(views.len() / 2);
34    views.iter().for_each(|view| {
35        distinct.insert(view.head());
36    });
37
38    Ok(u32::try_from(distinct.len())?)
39}
40
41impl StringStats {
42    /// Generates stats, returning an error on failure.
43    fn generate_opts_fallible(
44        input: &VarBinViewArray,
45        opts: GenerateStatsOptions,
46        ctx: &mut ExecutionCtx,
47    ) -> VortexResult<Self> {
48        let null_count = input
49            .statistics()
50            .compute_null_count(ctx)
51            .ok_or_else(|| vortex_err!("Failed to compute null_count"))?;
52        let value_count = input.len() - null_count;
53        let estimated_distinct_count = opts
54            .count_distinct_values
55            .then(|| estimate_distinct_count(input))
56            .transpose()?;
57
58        Ok(Self {
59            value_count: u32::try_from(value_count)?,
60            null_count: u32::try_from(null_count)?,
61            estimated_distinct_count,
62        })
63    }
64}
65
66impl StringStats {
67    /// Generates stats with default options.
68    pub fn generate(input: &VarBinViewArray, ctx: &mut ExecutionCtx) -> Self {
69        Self::generate_opts(input, GenerateStatsOptions::default(), ctx)
70    }
71
72    /// Generates stats with provided options.
73    pub fn generate_opts(
74        input: &VarBinViewArray,
75        opts: GenerateStatsOptions,
76        ctx: &mut ExecutionCtx,
77    ) -> Self {
78        Self::generate_opts_fallible(input, opts, ctx)
79            .vortex_expect("StringStats::generate_opts should not fail")
80    }
81
82    /// Returns the estimated number of distinct values, or `None` if not computed.
83    ///
84    /// This estimation is always going to be less than or equal to the actual distinct count.
85    pub fn estimated_distinct_count(&self) -> Option<u32> {
86        self.estimated_distinct_count
87    }
88
89    /// Returns the number of non-null values.
90    pub fn value_count(&self) -> u32 {
91        self.value_count
92    }
93
94    /// Returns the number of null values.
95    pub fn null_count(&self) -> u32 {
96        self.null_count
97    }
98}