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    #[cfg(feature = "unstable_encodings")]
41    &integer::DeltaScheme::new(1.25),
42    ////////////////////////////////////////////////////////////////////////////////////////////////
43    // Float schemes.
44    ////////////////////////////////////////////////////////////////////////////////////////////////
45    &float::ALPScheme,
46    &float::ALPRDScheme,
47    &float::FloatDictScheme,
48    &float::NullDominatedSparseScheme,
49    &float::FloatRLEScheme,
50    ////////////////////////////////////////////////////////////////////////////////////////////////
51    // String schemes.
52    ////////////////////////////////////////////////////////////////////////////////////////////////
53    &string::StringDictScheme,
54    // Both string-fragmentation schemes are registered; the sample-based
55    // selector keeps whichever is smaller per column.
56    &string::FSSTScheme,
57    #[cfg(feature = "unstable_encodings")]
58    &string::OnPairScheme,
59    &string::NullDominatedSparseScheme,
60    ////////////////////////////////////////////////////////////////////////////////////////////////
61    // Binary schemes.
62    ////////////////////////////////////////////////////////////////////////////////////////////////
63    &binary::BinaryDictScheme,
64    // Decimal schemes.
65    &decimal::DecimalScheme,
66    // Temporal schemes.
67    &temporal::TemporalScheme,
68];
69
70/// Builder for creating configured [`BtrBlocksCompressor`] instances.
71///
72/// By default, all schemes in [`ALL_SCHEMES`] are enabled in a deterministic order. Feature-gated
73/// schemes (Pco, Zstd) are not in `ALL_SCHEMES` and must be added explicitly via
74/// [`with_new_scheme`](BtrBlocksCompressorBuilder::with_new_scheme) or `with_compact` when the
75/// `zstd` feature is enabled.
76///
77/// # Examples
78///
79/// ```rust
80/// use vortex_btrblocks::{BtrBlocksCompressorBuilder, Scheme, SchemeExt};
81/// use vortex_btrblocks::schemes::integer::IntDictScheme;
82///
83/// // Default compressor with all schemes in ALL_SCHEMES.
84/// let compressor = BtrBlocksCompressorBuilder::default().build();
85///
86/// // Remove specific schemes.
87/// let compressor = BtrBlocksCompressorBuilder::default()
88///     .exclude_schemes([IntDictScheme.id()])
89///     .build();
90/// ```
91#[derive(Debug, Clone)]
92pub struct BtrBlocksCompressorBuilder {
93    schemes: Vec<&'static dyn Scheme>,
94}
95
96impl Default for BtrBlocksCompressorBuilder {
97    fn default() -> Self {
98        Self {
99            schemes: ALL_SCHEMES.to_vec(),
100        }
101    }
102}
103
104impl BtrBlocksCompressorBuilder {
105    /// Creates a builder with no schemes registered.
106    ///
107    /// Useful when the caller wants explicit, scheme-by-scheme control over the compressor.
108    pub fn empty() -> Self {
109        Self {
110            schemes: Vec::new(),
111        }
112    }
113
114    /// Adds an external compression scheme not in [`ALL_SCHEMES`].
115    ///
116    /// This allows encoding crates outside of `vortex-btrblocks` to register their own schemes
117    /// with the compressor.
118    ///
119    /// # Panics
120    ///
121    /// Panics if a scheme with the same [`SchemeId`] is already present.
122    pub fn with_new_scheme(mut self, scheme: &'static dyn Scheme) -> Self {
123        assert!(
124            !self.schemes.iter().any(|s| s.id() == scheme.id()),
125            "scheme {:?} is already present in the builder",
126            scheme.id(),
127        );
128
129        self.schemes.push(scheme);
130        self
131    }
132
133    /// Adds compact encoding schemes (Zstd for strings and binary, Pco for numerics).
134    ///
135    /// This provides better compression ratios than the default, especially for floating-point
136    /// heavy datasets. Requires the `zstd` feature. When the `pco` feature is also enabled,
137    /// Pco schemes for integers and floats are included.
138    ///
139    /// # Panics
140    ///
141    /// Panics if any of the compact schemes are already present.
142    #[cfg(feature = "zstd")]
143    pub fn with_compact(self) -> Self {
144        let builder = self
145            .with_new_scheme(&string::ZstdScheme)
146            .with_new_scheme(&binary::ZstdScheme);
147
148        #[cfg(feature = "pco")]
149        let builder = builder
150            .with_new_scheme(&integer::PcoScheme)
151            .with_new_scheme(&float::PcoScheme);
152
153        builder
154    }
155
156    /// Excludes schemes without CUDA kernel support, keeps FSST for string compression,
157    /// and adds Zstd for binary compression.
158    ///
159    /// With the `unstable_encodings` feature, buffer-level Zstd compression is used for binary
160    /// arrays, preserving their buffer layout for zero-conversion GPU decompression. Without it,
161    /// interleaved binary Zstd compression is used.
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(
169            not(any(feature = "pco", feature = "unstable_encodings")),
170            allow(unused_mut)
171        )]
172        let mut excluded: Vec<SchemeId> = vec![
173            integer::SparseScheme.id(),
174            integer::IntRLEScheme.id(),
175            float::ALPRDScheme.id(),
176            float::FloatRLEScheme.id(),
177            float::NullDominatedSparseScheme.id(),
178            string::StringDictScheme.id(),
179            binary::BinaryDictScheme.id(),
180        ];
181        #[cfg(feature = "unstable_encodings")]
182        excluded.push(string::OnPairScheme.id());
183        // Delta has no GPU decode kernel and its prefix-sum decode is inherently sequential, so it
184        // is incompatible with pure-GPU decompression paths.
185        #[cfg(feature = "unstable_encodings")]
186        excluded.push(integer::DeltaScheme::default().id());
187        #[cfg(feature = "pco")]
188        excluded.extend([integer::PcoScheme.id(), float::PcoScheme.id()]);
189        let builder = self.exclude_schemes(excluded);
190
191        #[cfg(all(feature = "zstd", feature = "unstable_encodings"))]
192        let builder = builder.with_new_scheme(&binary::ZstdBuffersScheme);
193        #[cfg(all(feature = "zstd", not(feature = "unstable_encodings")))]
194        let builder = builder.with_new_scheme(&binary::ZstdScheme);
195
196        builder
197    }
198
199    /// Removes the specified compression schemes by their [`SchemeId`].
200    pub fn exclude_schemes(mut self, ids: impl IntoIterator<Item = SchemeId>) -> Self {
201        let ids: HashSet<_> = ids.into_iter().collect();
202        self.schemes.retain(|s| !ids.contains(&s.id()));
203        self
204    }
205
206    /// Retains only schemes whose produced encodings all belong to `allowed`.
207    ///
208    /// The file writer uses this to restrict compression to the encodings of its configured
209    /// editions.
210    pub fn retain_allowed_encodings(mut self, allowed: &HashSet<ArrayId>) -> Self {
211        self.schemes
212            .retain(|s| s.produced_encodings().iter().all(|id| allowed.contains(id)));
213        self
214    }
215
216    /// Builds the configured [`BtrBlocksCompressor`].
217    pub fn build(self) -> BtrBlocksCompressor {
218        BtrBlocksCompressor(CascadingCompressor::new(self.schemes))
219    }
220}
221
222#[cfg(test)]
223mod tests {
224    use vortex_array::VTable;
225    use vortex_fastlanes::FoR;
226
227    use super::*;
228
229    #[test]
230    fn empty_starts_with_no_schemes() {
231        let builder = BtrBlocksCompressorBuilder::empty();
232        assert!(builder.schemes.is_empty());
233    }
234
235    #[test]
236    fn default_includes_all_schemes() {
237        let builder = BtrBlocksCompressorBuilder::default();
238        assert_eq!(builder.schemes.len(), ALL_SCHEMES.len());
239    }
240
241    #[test]
242    fn retain_allowed_encodings_filters_schemes() {
243        let allowed: HashSet<ArrayId> = [FoR.id()].into_iter().collect();
244        let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed);
245        assert_eq!(builder.schemes.len(), 1);
246        assert_eq!(builder.schemes[0].id(), integer::FoRScheme.id());
247
248        let none = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&HashSet::new());
249        assert!(none.schemes.is_empty());
250    }
251
252    #[test]
253    fn retaining_all_declared_outputs_keeps_every_scheme() {
254        let allowed: HashSet<ArrayId> = ALL_SCHEMES
255            .iter()
256            .flat_map(|scheme| scheme.produced_encodings())
257            .collect();
258        let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed);
259        assert_eq!(builder.schemes.len(), ALL_SCHEMES.len());
260    }
261
262    #[test]
263    fn cuda_compatible_excludes_alprd() {
264        let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible();
265        assert!(
266            !builder
267                .schemes
268                .iter()
269                .any(|s| s.id() == float::ALPRDScheme.id())
270        );
271    }
272
273    #[test]
274    fn cuda_compatible_uses_fsst_for_strings() {
275        let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible();
276        assert!(
277            builder
278                .schemes
279                .iter()
280                .any(|scheme| scheme.id() == string::FSSTScheme.id())
281        );
282        #[cfg(feature = "zstd")]
283        assert!(
284            !builder
285                .schemes
286                .iter()
287                .any(|scheme| scheme.id() == string::ZstdScheme.id())
288        );
289    }
290
291    #[test]
292    #[cfg(feature = "pco")]
293    fn cuda_compatible_excludes_pco() {
294        let builder = BtrBlocksCompressorBuilder::default()
295            .with_new_scheme(&integer::PcoScheme)
296            .with_new_scheme(&float::PcoScheme)
297            .only_cuda_compatible();
298        for scheme in [integer::PcoScheme.id(), float::PcoScheme.id()] {
299            assert!(!builder.schemes.iter().any(|s| s.id() == scheme));
300        }
301    }
302}
vortex_btrblocks/builder.rs

vortex_btrblocks/
builder.rs
