vortex_btrblocks/schemes/integer/
sparse.rs1use 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#[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 fn num_children(&self) -> usize {
61 2
62 }
63
64 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 if value_count == 0 {
99 return CompressionEstimate::Verdict(EstimateVerdict::Skip);
100 }
101
102 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 most_frequent_count == value_count {
116 return CompressionEstimate::Verdict(EstimateVerdict::Skip);
117 }
118 debug_assert!(value_count > most_frequent_count);
119
120 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 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 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}