Skip to main content

vortex_btrblocks/schemes/integer/
sparse.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Sparse integer encoding for single-value-dominated arrays.
5
6use vortex_array::ArrayId;
7use vortex_array::ArrayRef;
8use vortex_array::Canonical;
9use vortex_array::ExecutionCtx;
10use vortex_array::IntoArray;
11use vortex_array::VTable;
12use vortex_array::arrays::Constant;
13use vortex_array::arrays::ConstantArray;
14use vortex_array::arrays::PrimitiveArray;
15use vortex_array::arrays::primitive::PrimitiveArrayExt;
16use vortex_array::scalar::Scalar;
17use vortex_compressor::builtins::IntDictScheme;
18use vortex_compressor::scheme::ChildSelection;
19use vortex_compressor::scheme::CompressionEstimate;
20use vortex_compressor::scheme::DescendantExclusion;
21use vortex_compressor::scheme::EstimateVerdict;
22use vortex_error::VortexExpect;
23use vortex_error::VortexResult;
24use vortex_sparse::Sparse;
25use vortex_sparse::SparseExt as _;
26
27use super::IntRLEScheme;
28use super::RunEndScheme;
29use crate::ArrayAndStats;
30use crate::CascadingCompressor;
31use crate::CompressorContext;
32use crate::GenerateStatsOptions;
33use crate::Scheme;
34use crate::SchemeExt;
35
36/// Sparse encoding for single-value-dominated arrays.
37#[derive(Debug, Copy, Clone, PartialEq, Eq)]
38pub struct SparseScheme;
39
40impl Scheme for SparseScheme {
41    fn scheme_name(&self) -> &'static str {
42        "vortex.int.sparse"
43    }
44
45    fn matches(&self, canonical: &Canonical) -> bool {
46        canonical.dtype().is_int()
47    }
48
49    fn produced_encodings(&self) -> Vec<ArrayId> {
50        vec![Sparse.id(), Constant.id()]
51    }
52
53    fn stats_options(&self) -> GenerateStatsOptions {
54        GenerateStatsOptions {
55            count_distinct_values: true,
56        }
57    }
58
59    /// Children: values=0, indices=1.
60    fn num_children(&self) -> usize {
61        2
62    }
63
64    /// Sparse indices (child 1) are monotonically increasing positions with all unique values.
65    /// Dict, RunEnd, RLE, and Sparse are all pointless on such data.
66    fn descendant_exclusions(&self) -> Vec<DescendantExclusion> {
67        vec![
68            DescendantExclusion {
69                excluded: IntDictScheme.id(),
70                children: ChildSelection::One(1),
71            },
72            DescendantExclusion {
73                excluded: RunEndScheme.id(),
74                children: ChildSelection::One(1),
75            },
76            DescendantExclusion {
77                excluded: IntRLEScheme.id(),
78                children: ChildSelection::One(1),
79            },
80            DescendantExclusion {
81                excluded: SparseScheme.id(),
82                children: ChildSelection::One(1),
83            },
84        ]
85    }
86
87    fn expected_compression_ratio(
88        &self,
89        data: &ArrayAndStats,
90        _compress_ctx: CompressorContext,
91        exec_ctx: &mut ExecutionCtx,
92    ) -> CompressionEstimate {
93        let len = data.array_len() as f64;
94        let stats = data.integer_stats(exec_ctx);
95        let value_count = stats.value_count();
96
97        // All-null arrays should be compressed as constant instead anyways.
98        if value_count == 0 {
99            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
100        }
101
102        // If the majority (90%) of values is null, this will compress well.
103        if stats.null_count() as f64 / len > 0.9 {
104            return CompressionEstimate::Verdict(EstimateVerdict::Ratio(len / value_count as f64));
105        }
106
107        let (_, most_frequent_count) = stats
108            .erased()
109            .most_frequent_value_and_count()
110            .vortex_expect(
111                "this must be present since `SparseScheme` declared that we need distinct values",
112            );
113
114        // If the most frequent value is the only value, we should compress as constant instead.
115        if most_frequent_count == value_count {
116            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
117        }
118        debug_assert!(value_count > most_frequent_count);
119
120        // See if the most frequent value accounts for >= 90% of the set values.
121        let freq = most_frequent_count as f64 / value_count as f64;
122        if freq < 0.9 {
123            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
124        }
125
126        // We only store the positions of the non-top values.
127        CompressionEstimate::Verdict(EstimateVerdict::Ratio(
128            value_count as f64 / (value_count - most_frequent_count) as f64,
129        ))
130    }
131
132    fn compress(
133        &self,
134        compressor: &CascadingCompressor,
135        data: &ArrayAndStats,
136        compress_ctx: CompressorContext,
137        exec_ctx: &mut ExecutionCtx,
138    ) -> VortexResult<ArrayRef> {
139        let len = data.array_len();
140        let stats = data.integer_stats(exec_ctx);
141        let array = data.array();
142
143        let (most_frequent_value, most_frequent_count) = stats
144            .erased()
145            .most_frequent_value_and_count()
146            .vortex_expect(
147                "this must be present since `SparseScheme` declared that we need distinct values",
148            );
149
150        if most_frequent_count as usize == len {
151            // If the most frequent value is the only value, we should compress as constant instead.
152            return Ok(ConstantArray::new(
153                Scalar::primitive_value(
154                    most_frequent_value,
155                    most_frequent_value.ptype(),
156                    array.dtype().nullability(),
157                ),
158                len,
159            )
160            .into_array());
161        }
162
163        let sparse_encoded = Sparse::encode(
164            array,
165            Some(Scalar::primitive_value(
166                most_frequent_value,
167                most_frequent_value.ptype(),
168                array.dtype().nullability(),
169            )),
170            exec_ctx,
171        )?;
172
173        if let Some(sparse) = sparse_encoded.as_opt::<Sparse>() {
174            let sparse_values_primitive = sparse
175                .patches()
176                .values()
177                .clone()
178                .execute::<PrimitiveArray>(exec_ctx)?;
179            let compressed_values = compressor.compress_child(
180                &sparse_values_primitive.into_array(),
181                &compress_ctx,
182                self.id(),
183                0,
184                exec_ctx,
185            )?;
186
187            let indices = sparse
188                .patches()
189                .indices()
190                .clone()
191                .execute::<PrimitiveArray>(exec_ctx)?
192                .narrow(exec_ctx)?;
193
194            let compressed_indices = compressor.compress_child(
195                &indices.into_array(),
196                &compress_ctx,
197                self.id(),
198                1,
199                exec_ctx,
200            )?;
201
202            Sparse::try_new(
203                compressed_indices,
204                compressed_values,
205                sparse.len(),
206                sparse.fill_scalar().clone(),
207            )
208            .map(|a| a.into_array())
209        } else {
210            Ok(sparse_encoded)
211        }
212    }
213}