Skip to main content

vortex_btrblocks/schemes/float/
sparse.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Sparse encoding for null-dominated float 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::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/// Sparse encoding for null-dominated float arrays.
30///
31/// This is the same as the integer `SparseScheme`, but we only use this for null-dominated arrays.
32#[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    /// Children: indices=0.
49    fn num_children(&self) -> usize {
50        1
51    }
52
53    /// The indices of a null-dominated sparse array should not be sparse-encoded again.
54    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        // All-null arrays should be compressed as constant instead anyways.
72        if value_count == 0 {
73            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
74        }
75
76        // If the majority (90%) of values is null, this will compress well.
77        if stats.null_count() as f64 / len > 0.9 {
78            return CompressionEstimate::Verdict(EstimateVerdict::Ratio(len / value_count as f64));
79        }
80
81        // Otherwise we don't go this route.
82        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        // We pass None as we only run this pathway for NULL-dominated float arrays.
93        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}