vortex_btrblocks/
builder.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Builder for configuring `BtrBlocksCompressor` instances.
5
6use vortex_array::ArrayId;
7use vortex_utils::aliases::hash_set::HashSet;
8
9use crate::BtrBlocksCompressor;
10use crate::CascadingCompressor;
11use crate::Scheme;
12use crate::SchemeExt;
13use crate::SchemeId;
14use crate::schemes::binary;
15use crate::schemes::decimal;
16use crate::schemes::float;
17use crate::schemes::integer;
18use crate::schemes::string;
19use crate::schemes::temporal;
20
21/// All available compression schemes.
22///
23/// This list is order-sensitive: the builder preserves this order when constructing
24/// the final scheme list, so that tie-breaking is deterministic.
25pub const ALL_SCHEMES: &[&dyn Scheme] = &[
26    ////////////////////////////////////////////////////////////////////////////////////////////////
27    // Integer schemes.
28    ////////////////////////////////////////////////////////////////////////////////////////////////
29    // NOTE: FoR must precede BitPacking to avoid unnecessary patches.
30    &integer::FoRScheme,
31    // NOTE: ZigZag should precede BitPacking because we don't want negative numbers.
32    &integer::ZigZagScheme,
33    &integer::BitPackingScheme,
34    &integer::SparseScheme,
35    &integer::IntDictScheme,
36    &integer::RunEndScheme,
37    &integer::SequenceScheme,
38    &integer::IntRLEScheme,
39    // Prefer all other schemes above delta, for now (since its slower to decompress).
40    &integer::DeltaScheme::new(1.25),
41    ////////////////////////////////////////////////////////////////////////////////////////////////
42    // Float schemes.
43    ////////////////////////////////////////////////////////////////////////////////////////////////
44    &float::ALPScheme,
45    &float::ALPRDScheme,
46    &float::FloatDictScheme,
47    &float::NullDominatedSparseScheme,
48    &float::FloatRLEScheme,
49    ////////////////////////////////////////////////////////////////////////////////////////////////
50    // String schemes.
51    ////////////////////////////////////////////////////////////////////////////////////////////////
52    &string::StringDictScheme,
53    // Both string-fragmentation schemes are registered; the sample-based
54    // selector keeps whichever is smaller per column.
55    &string::FSSTScheme,
56    &string::OnPairScheme,
57    &string::NullDominatedSparseScheme,
58    ////////////////////////////////////////////////////////////////////////////////////////////////
59    // Binary schemes.
60    ////////////////////////////////////////////////////////////////////////////////////////////////
61    &binary::BinaryDictScheme,
62    &binary::VarBinScheme,
63    // Decimal schemes.
64    &decimal::DecimalScheme,
65    // Temporal schemes.
66    &temporal::TemporalScheme,
67];
68
69/// Builder for creating configured [`BtrBlocksCompressor`] instances.
70///
71/// By default, all schemes in [`ALL_SCHEMES`] are enabled in a deterministic order. Feature-gated
72/// schemes (Pco, Zstd) are not in `ALL_SCHEMES` and must be added explicitly via
73/// [`with_new_scheme`](BtrBlocksCompressorBuilder::with_new_scheme) or `with_compact` when the
74/// `zstd` feature is enabled.
75///
76/// # Examples
77///
78/// ```rust
79/// use vortex_btrblocks::{BtrBlocksCompressorBuilder, Scheme, SchemeExt};
80/// use vortex_btrblocks::schemes::integer::IntDictScheme;
81///
82/// // Default compressor with all schemes in ALL_SCHEMES.
83/// let compressor = BtrBlocksCompressorBuilder::default().build();
84///
85/// // Remove specific schemes.
86/// let compressor = BtrBlocksCompressorBuilder::default()
87///     .exclude_schemes([IntDictScheme.id()])
88///     .build();
89/// ```
90#[derive(Debug, Clone)]
91pub struct BtrBlocksCompressorBuilder {
92    schemes: Vec<&'static dyn Scheme>,
93}
94
95impl Default for BtrBlocksCompressorBuilder {
96    fn default() -> Self {
97        Self {
98            schemes: ALL_SCHEMES.to_vec(),
99        }
100    }
101}
102
103impl BtrBlocksCompressorBuilder {
104    /// Creates a builder with no schemes registered.
105    ///
106    /// Useful when the caller wants explicit, scheme-by-scheme control over the compressor.
107    pub fn empty() -> Self {
108        Self {
109            schemes: Vec::new(),
110        }
111    }
112
113    /// Adds an external compression scheme not in [`ALL_SCHEMES`].
114    ///
115    /// This allows encoding crates outside of `vortex-btrblocks` to register their own schemes
116    /// with the compressor.
117    ///
118    /// # Panics
119    ///
120    /// Panics if a scheme with the same [`SchemeId`] is already present.
121    pub fn with_new_scheme(mut self, scheme: &'static dyn Scheme) -> Self {
122        assert!(
123            !self.schemes.iter().any(|s| s.id() == scheme.id()),
124            "scheme {:?} is already present in the builder",
125            scheme.id(),
126        );
127
128        self.schemes.push(scheme);
129        self
130    }
131
132    /// Adds compact encoding schemes (Zstd for strings and binary, Pco for numerics).
133    ///
134    /// This provides better compression ratios than the default, especially for floating-point
135    /// heavy datasets. Requires the `zstd` feature. When the `pco` feature is also enabled,
136    /// Pco schemes for integers and floats are included.
137    ///
138    /// # Panics
139    ///
140    /// Panics if any of the compact schemes are already present.
141    #[cfg(feature = "zstd")]
142    pub fn with_compact(self) -> Self {
143        let builder = self
144            .with_new_scheme(&string::ZstdScheme)
145            .with_new_scheme(&binary::ZstdScheme);
146
147        #[cfg(feature = "pco")]
148        let builder = builder
149            .with_new_scheme(&integer::PcoScheme)
150            .with_new_scheme(&float::PcoScheme);
151
152        builder
153    }
154
155    /// Excludes schemes without CUDA kernel support, keeps FSST for string compression,
156    /// and adds Zstd for binary compression.
157    ///
158    /// Both the array-level and the buffer-level Zstd schemes are added. Buffer-level
159    /// compression preserves binary arrays' buffer layout for zero-conversion GPU decompression,
160    /// but belongs to the opt-in `zstd` edition, so callers filter the two through
161    /// [`retain_allowed_encodings`](Self::retain_allowed_encodings).
162    ///
163    /// This preset is intended for files that will be decoded by CUDA kernels. It may choose a
164    /// larger encoded representation than the default compressor.
165    pub fn only_cuda_compatible(self) -> Self {
166        // Keep FSST, which has a CUDA decoder and direct Arrow offset-based export. Other
167        // string fragmentation and dictionary schemes still require unsupported decode paths.
168        #[cfg_attr(not(any(feature = "pco", feature = "zstd")), allow(unused_mut))]
169        let mut excluded: Vec<SchemeId> = vec![
170            integer::SparseScheme.id(),
171            integer::IntRLEScheme.id(),
172            float::ALPRDScheme.id(),
173            float::FloatRLEScheme.id(),
174            float::NullDominatedSparseScheme.id(),
175            string::NullDominatedSparseScheme.id(),
176            string::StringDictScheme.id(),
177            binary::BinaryDictScheme.id(),
178        ];
179        // Delta now has a CUDA decode kernel, so arrays that reach the GPU already encoded with
180        // it — the Delta children OnPair emits, for instance — decode there. It stays excluded
181        // from this preset until GPU delta decode is benchmarked against the schemes it would
182        // displace, since the preset picks encodings rather than merely decoding them.
183        excluded.push(integer::DeltaScheme::default().id());
184        #[cfg(feature = "pco")]
185        excluded.extend([integer::PcoScheme.id(), float::PcoScheme.id()]);
186        let builder = self.exclude_schemes(excluded);
187
188        #[cfg(feature = "zstd")]
189        let builder = builder
190            .with_new_scheme(&binary::ZstdScheme)
191            .with_new_scheme(&binary::ZstdBuffersScheme);
192
193        builder
194    }
195
196    /// Removes the specified compression schemes by their [`SchemeId`].
197    pub fn exclude_schemes(mut self, ids: impl IntoIterator<Item = SchemeId>) -> Self {
198        let ids: HashSet<_> = ids.into_iter().collect();
199        self.schemes.retain(|s| !ids.contains(&s.id()));
200        self
201    }
202
203    /// Retains only schemes whose produced encodings all belong to `allowed`.
204    ///
205    /// The file writer uses this to restrict compression to the encodings of its configured
206    /// editions.
207    pub fn retain_allowed_encodings(mut self, allowed: &HashSet<ArrayId>) -> Self {
208        self.schemes
209            .retain(|s| s.produced_encodings().iter().all(|id| allowed.contains(id)));
210        self
211    }
212
213    /// Builds the configured [`BtrBlocksCompressor`].
214    pub fn build(self) -> BtrBlocksCompressor {
215        BtrBlocksCompressor(CascadingCompressor::new(self.schemes))
216    }
217}
218
219#[cfg(test)]
220mod tests {
221    use vortex_array::VTable;
222    use vortex_fastlanes::FoR;
223
224    use super::*;
225
226    #[test]
227    fn empty_starts_with_no_schemes() {
228        let builder = BtrBlocksCompressorBuilder::empty();
229        assert!(builder.schemes.is_empty());
230    }
231
232    #[test]
233    fn default_includes_all_schemes() {
234        let builder = BtrBlocksCompressorBuilder::default();
235        assert_eq!(builder.schemes.len(), ALL_SCHEMES.len());
236    }
237
238    #[test]
239    fn retain_allowed_encodings_filters_schemes() {
240        let allowed: HashSet<ArrayId> = [FoR.id()].into_iter().collect();
241        let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed);
242        assert_eq!(builder.schemes.len(), 1);
243        assert_eq!(builder.schemes[0].id(), integer::FoRScheme.id());
244
245        let none = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&HashSet::new());
246        assert!(none.schemes.is_empty());
247    }
248
249    #[test]
250    fn retaining_all_declared_outputs_keeps_every_scheme() {
251        let allowed: HashSet<ArrayId> = ALL_SCHEMES
252            .iter()
253            .flat_map(|scheme| scheme.produced_encodings())
254            .collect();
255        let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed);
256        assert_eq!(builder.schemes.len(), ALL_SCHEMES.len());
257    }
258
259    #[test]
260    fn cuda_compatible_excludes_alprd() {
261        let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible();
262        assert!(
263            !builder
264                .schemes
265                .iter()
266                .any(|s| s.id() == float::ALPRDScheme.id())
267        );
268    }
269
270    /// `vortex.sparse` has no CUDA decode kernel, so no sparse scheme may survive this preset.
271    #[test]
272    fn cuda_compatible_excludes_every_sparse_scheme() {
273        let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible();
274        for excluded in [
275            integer::SparseScheme.id(),
276            float::NullDominatedSparseScheme.id(),
277            string::NullDominatedSparseScheme.id(),
278        ] {
279            assert!(
280                !builder.schemes.iter().any(|s| s.id() == excluded),
281                "{excluded} should be excluded"
282            );
283        }
284    }
285
286    #[test]
287    fn cuda_compatible_uses_fsst_for_strings() {
288        let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible();
289        assert!(
290            builder
291                .schemes
292                .iter()
293                .any(|scheme| scheme.id() == string::FSSTScheme.id())
294        );
295        #[cfg(feature = "zstd")]
296        assert!(
297            !builder
298                .schemes
299                .iter()
300                .any(|scheme| scheme.id() == string::ZstdScheme.id())
301        );
302    }
303
304    #[test]
305    #[cfg(feature = "pco")]
306    fn cuda_compatible_excludes_pco() {
307        let builder = BtrBlocksCompressorBuilder::default()
308            .with_new_scheme(&integer::PcoScheme)
309            .with_new_scheme(&float::PcoScheme)
310            .only_cuda_compatible();
311        for scheme in [integer::PcoScheme.id(), float::PcoScheme.id()] {
312            assert!(!builder.schemes.iter().any(|s| s.id() == scheme));
313        }
314    }
315}
vortex_btrblocks/builder.rs

vortex_btrblocks/
builder.rs
