vortex_btrblocks/schemes/float/
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::PrimitiveArray;
13use vortex_array::arrays::primitive::PrimitiveArrayExt;
14use vortex_compressor::scheme::ChildSelection;
15use vortex_compressor::scheme::CompressionEstimate;
16use vortex_compressor::scheme::DescendantExclusion;
17use vortex_compressor::scheme::EstimateVerdict;
18use vortex_error::VortexResult;
19use vortex_sparse::Sparse;
20use vortex_sparse::SparseExt as _;
21
22use crate::ArrayAndStats;
23use crate::CascadingCompressor;
24use crate::CompressorContext;
25use crate::Scheme;
26use crate::SchemeExt;
27use crate::schemes::integer::SparseScheme as IntSparseScheme;
28
29#[derive(Debug, Copy, Clone, PartialEq, Eq)]
33pub struct NullDominatedSparseScheme;
34
35impl Scheme for NullDominatedSparseScheme {
36 fn scheme_name(&self) -> &'static str {
37 "vortex.float.sparse"
38 }
39
40 fn matches(&self, canonical: &Canonical) -> bool {
41 canonical.dtype().is_float()
42 }
43
44 fn produced_encodings(&self) -> Vec<ArrayId> {
45 vec![Sparse.id()]
46 }
47
48 fn num_children(&self) -> usize {
50 1
51 }
52
53 fn descendant_exclusions(&self) -> Vec<DescendantExclusion> {
55 vec![DescendantExclusion {
56 excluded: IntSparseScheme.id(),
57 children: ChildSelection::All,
58 }]
59 }
60
61 fn expected_compression_ratio(
62 &self,
63 data: &ArrayAndStats,
64 _compress_ctx: CompressorContext,
65 exec_ctx: &mut ExecutionCtx,
66 ) -> CompressionEstimate {
67 let len = data.array_len() as f64;
68 let stats = data.float_stats(exec_ctx);
69 let value_count = stats.value_count();
70
71 if value_count == 0 {
73 return CompressionEstimate::Verdict(EstimateVerdict::Skip);
74 }
75
76 if stats.null_count() as f64 / len > 0.9 {
78 return CompressionEstimate::Verdict(EstimateVerdict::Ratio(len / value_count as f64));
79 }
80
81 CompressionEstimate::Verdict(EstimateVerdict::Skip)
83 }
84
85 fn compress(
86 &self,
87 compressor: &CascadingCompressor,
88 data: &ArrayAndStats,
89 compress_ctx: CompressorContext,
90 exec_ctx: &mut ExecutionCtx,
91 ) -> VortexResult<ArrayRef> {
92 let sparse_encoded = Sparse::encode(data.array(), None, exec_ctx)?;
94
95 if let Some(sparse) = sparse_encoded.as_opt::<Sparse>() {
96 let indices = sparse
97 .patches()
98 .indices()
99 .clone()
100 .execute::<PrimitiveArray>(exec_ctx)?
101 .narrow(exec_ctx)?;
102 let compressed_indices = compressor.compress_child(
103 &indices.into_array(),
104 &compress_ctx,
105 self.id(),
106 0,
107 exec_ctx,
108 )?;
109
110 Sparse::try_new(
111 compressed_indices,
112 sparse.patches().values().clone(),
113 sparse.len(),
114 sparse.fill_scalar().clone(),
115 )
116 .map(|a| a.into_array())
117 } else {
118 Ok(sparse_encoded)
119 }
120 }
121}