Skip to main content

vortex_btrblocks/schemes/string/
sparse.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Sparse encoding for null-dominated string 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::IntDictScheme;
28use crate::schemes::integer::SparseScheme as IntSparseScheme;
29
30/// Sparse encoding for null-dominated arrays.
31///
32/// This is the same as the integer `SparseScheme`, but we only use this for null-dominated arrays.
33#[derive(Debug, Copy, Clone, PartialEq, Eq)]
34pub struct NullDominatedSparseScheme;
35
36impl Scheme for NullDominatedSparseScheme {
37    fn scheme_name(&self) -> &'static str {
38        "vortex.string.sparse"
39    }
40
41    fn matches(&self, canonical: &Canonical) -> bool {
42        canonical.dtype().is_utf8()
43    }
44
45    fn produced_encodings(&self) -> Vec<ArrayId> {
46        vec![Sparse.id()]
47    }
48
49    /// Children: indices=0.
50    fn num_children(&self) -> usize {
51        1
52    }
53
54    /// The indices of a null-dominated sparse array should not be sparse-encoded again.
55    fn descendant_exclusions(&self) -> Vec<DescendantExclusion> {
56        vec![
57            DescendantExclusion {
58                excluded: IntSparseScheme.id(),
59                children: ChildSelection::All,
60            },
61            DescendantExclusion {
62                excluded: IntDictScheme.id(),
63                children: ChildSelection::All,
64            },
65        ]
66    }
67
68    fn expected_compression_ratio(
69        &self,
70        data: &ArrayAndStats,
71        _compress_ctx: CompressorContext,
72        exec_ctx: &mut ExecutionCtx,
73    ) -> CompressionEstimate {
74        let len = data.array_len() as f64;
75        let stats = data.varbinview_stats(exec_ctx);
76        let value_count = stats.value_count();
77
78        // All-null arrays should be compressed as constant instead anyways.
79        if value_count == 0 {
80            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
81        }
82
83        // If the majority (90%) of values is null, this will compress well.
84        if stats.null_count() as f64 / len > 0.9 {
85            return CompressionEstimate::Verdict(EstimateVerdict::Ratio(len / value_count as f64));
86        }
87
88        // Otherwise we don't go this route.
89        CompressionEstimate::Verdict(EstimateVerdict::Skip)
90    }
91
92    fn compress(
93        &self,
94        compressor: &CascadingCompressor,
95        data: &ArrayAndStats,
96        compress_ctx: CompressorContext,
97        exec_ctx: &mut ExecutionCtx,
98    ) -> VortexResult<ArrayRef> {
99        // We pass None as we only run this pathway for NULL-dominated string arrays.
100        let sparse_encoded = Sparse::encode(data.array(), None, exec_ctx)?;
101
102        if let Some(sparse) = sparse_encoded.as_opt::<Sparse>() {
103            // Compress the indices only (not the values for strings).
104            let indices = sparse
105                .patches()
106                .indices()
107                .clone()
108                .execute::<PrimitiveArray>(exec_ctx)?
109                .narrow(exec_ctx)?;
110            let compressed_indices = compressor.compress_child(
111                &indices.into_array(),
112                &compress_ctx,
113                self.id(),
114                0,
115                exec_ctx,
116            )?;
117
118            Sparse::try_new(
119                compressed_indices,
120                sparse.patches().values().clone(),
121                sparse.len(),
122                sparse.fill_scalar().clone(),
123            )
124            .map(|a| a.into_array())
125        } else {
126            Ok(sparse_encoded)
127        }
128    }
129}