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_utils::aliases::hash_set::HashSet;
7
8use crate::BtrBlocksCompressor;
9use crate::CascadingCompressor;
10use crate::Scheme;
11use crate::SchemeExt;
12use crate::SchemeId;
13use crate::schemes::binary;
14use crate::schemes::decimal;
15use crate::schemes::float;
16use crate::schemes::integer;
17use crate::schemes::string;
18use crate::schemes::temporal;
19
20/// All available compression schemes.
21///
22/// This list is order-sensitive: the builder preserves this order when constructing
23/// the final scheme list, so that tie-breaking is deterministic.
24pub const ALL_SCHEMES: &[&dyn Scheme] = &[
25    ////////////////////////////////////////////////////////////////////////////////////////////////
26    // Integer schemes.
27    ////////////////////////////////////////////////////////////////////////////////////////////////
28    // NOTE: FoR must precede BitPacking to avoid unnecessary patches.
29    &integer::FoRScheme,
30    // NOTE: ZigZag should precede BitPacking because we don't want negative numbers.
31    &integer::ZigZagScheme,
32    &integer::BitPackingScheme,
33    &integer::SparseScheme,
34    &integer::IntDictScheme,
35    &integer::RunEndScheme,
36    &integer::SequenceScheme,
37    &integer::IntRLEScheme,
38    // Prefer all other schemes above delta, for now (since its slower to decompress).
39    #[cfg(feature = "unstable_encodings")]
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    #[cfg(feature = "unstable_encodings")]
57    &string::OnPairScheme,
58    &string::NullDominatedSparseScheme,
59    ////////////////////////////////////////////////////////////////////////////////////////////////
60    // Binary schemes.
61    ////////////////////////////////////////////////////////////////////////////////////////////////
62    &binary::BinaryDictScheme,
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 and adds Zstd for string and binary compression.
156    ///
157    /// With the `unstable_encodings` feature, buffer-level Zstd compression is used which
158    /// preserves the array buffer layout for zero-conversion GPU decompression. Without it,
159    /// interleaved Zstd compression is used.
160    ///
161    /// This preset is intended for files that will be decoded by CUDA kernels. It may choose a
162    /// larger encoded representation than the default compressor.
163    pub fn only_cuda_compatible(self) -> Self {
164        // String fragmentation schemes (OnPair, FSST) require host-side
165        // dictionary expansion at decode time, which is incompatible with
166        // pure-GPU decompression paths. Strip whichever string-fragment
167        // scheme is enabled by feature.
168        #[cfg_attr(not(feature = "unstable_encodings"), allow(unused_mut))]
169        let mut excluded: Vec<SchemeId> = vec![
170            integer::SparseScheme.id(),
171            integer::IntRLEScheme.id(),
172            float::FloatRLEScheme.id(),
173            float::NullDominatedSparseScheme.id(),
174            string::StringDictScheme.id(),
175            string::FSSTScheme.id(),
176            binary::BinaryDictScheme.id(),
177        ];
178        #[cfg(feature = "unstable_encodings")]
179        excluded.push(string::OnPairScheme.id());
180        // Delta has no GPU decode kernel and its prefix-sum decode is inherently sequential, so it
181        // is incompatible with pure-GPU decompression paths.
182        #[cfg(feature = "unstable_encodings")]
183        excluded.push(integer::DeltaScheme::default().id());
184        let builder = self.exclude_schemes(excluded);
185
186        #[cfg(all(feature = "zstd", feature = "unstable_encodings"))]
187        let builder = builder
188            .with_new_scheme(&string::ZstdBuffersScheme)
189            .with_new_scheme(&binary::ZstdBuffersScheme);
190        #[cfg(all(feature = "zstd", not(feature = "unstable_encodings")))]
191        let builder = builder
192            .with_new_scheme(&string::ZstdScheme)
193            .with_new_scheme(&binary::ZstdScheme);
194
195        builder
196    }
197
198    /// Removes the specified compression schemes by their [`SchemeId`].
199    pub fn exclude_schemes(mut self, ids: impl IntoIterator<Item = SchemeId>) -> Self {
200        let ids: HashSet<_> = ids.into_iter().collect();
201        self.schemes.retain(|s| !ids.contains(&s.id()));
202        self
203    }
204
205    /// Builds the configured [`BtrBlocksCompressor`].
206    pub fn build(self) -> BtrBlocksCompressor {
207        BtrBlocksCompressor(CascadingCompressor::new(self.schemes))
208    }
209}
210
211#[cfg(test)]
212mod tests {
213    use super::*;
214
215    #[test]
216    fn empty_starts_with_no_schemes() {
217        let builder = BtrBlocksCompressorBuilder::empty();
218        assert!(builder.schemes.is_empty());
219    }
220
221    #[test]
222    fn default_includes_all_schemes() {
223        let builder = BtrBlocksCompressorBuilder::default();
224        assert_eq!(builder.schemes.len(), ALL_SCHEMES.len());
225    }
226}
vortex_btrblocks/builder.rs

vortex_btrblocks/
builder.rs
