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::bool;
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    // Bool schemes.
28    ////////////////////////////////////////////////////////////////////////////////////////////////
29    &bool::BoolConstantScheme,
30    ////////////////////////////////////////////////////////////////////////////////////////////////
31    // Integer schemes.
32    ////////////////////////////////////////////////////////////////////////////////////////////////
33    &integer::IntConstantScheme,
34    // NOTE: FoR must precede BitPacking to avoid unnecessary patches.
35    &integer::FoRScheme,
36    // NOTE: ZigZag should precede BitPacking because we don't want negative numbers.
37    &integer::ZigZagScheme,
38    &integer::BitPackingScheme,
39    &integer::SparseScheme,
40    &integer::IntDictScheme,
41    &integer::RunEndScheme,
42    &integer::SequenceScheme,
43    &integer::IntRLEScheme,
44    // Prefer all other schemes above delta, for now (since its slower to decompress).
45    #[cfg(feature = "unstable_encodings")]
46    &integer::DeltaScheme::new(1.25),
47    ////////////////////////////////////////////////////////////////////////////////////////////////
48    // Float schemes.
49    ////////////////////////////////////////////////////////////////////////////////////////////////
50    &float::FloatConstantScheme,
51    &float::ALPScheme,
52    &float::ALPRDScheme,
53    &float::FloatDictScheme,
54    &float::NullDominatedSparseScheme,
55    &float::FloatRLEScheme,
56    ////////////////////////////////////////////////////////////////////////////////////////////////
57    // String schemes.
58    ////////////////////////////////////////////////////////////////////////////////////////////////
59    &string::StringDictScheme,
60    // Both string-fragmentation schemes are registered; the sample-based
61    // selector keeps whichever is smaller per column.
62    &string::FSSTScheme,
63    #[cfg(feature = "unstable_encodings")]
64    &string::OnPairScheme,
65    &string::StringConstantScheme,
66    &string::NullDominatedSparseScheme,
67    ////////////////////////////////////////////////////////////////////////////////////////////////
68    // Binary schemes.
69    ////////////////////////////////////////////////////////////////////////////////////////////////
70    &binary::BinaryDictScheme,
71    &binary::BinaryConstantScheme,
72    // Decimal schemes.
73    &decimal::DecimalScheme,
74    // Temporal schemes.
75    &temporal::TemporalScheme,
76];
77
78/// Builder for creating configured [`BtrBlocksCompressor`] instances.
79///
80/// By default, all schemes in [`ALL_SCHEMES`] are enabled in a deterministic order. Feature-gated
81/// schemes (Pco, Zstd) are not in `ALL_SCHEMES` and must be added explicitly via
82/// [`with_new_scheme`](BtrBlocksCompressorBuilder::with_new_scheme) or `with_compact` when the
83/// `zstd` feature is enabled.
84///
85/// # Examples
86///
87/// ```rust
88/// use vortex_btrblocks::{BtrBlocksCompressorBuilder, Scheme, SchemeExt};
89/// use vortex_btrblocks::schemes::integer::IntDictScheme;
90///
91/// // Default compressor with all schemes in ALL_SCHEMES.
92/// let compressor = BtrBlocksCompressorBuilder::default().build();
93///
94/// // Remove specific schemes.
95/// let compressor = BtrBlocksCompressorBuilder::default()
96///     .exclude_schemes([IntDictScheme.id()])
97///     .build();
98/// ```
99#[derive(Debug, Clone)]
100pub struct BtrBlocksCompressorBuilder {
101    schemes: Vec<&'static dyn Scheme>,
102}
103
104impl Default for BtrBlocksCompressorBuilder {
105    fn default() -> Self {
106        Self {
107            schemes: ALL_SCHEMES.to_vec(),
108        }
109    }
110}
111
112impl BtrBlocksCompressorBuilder {
113    /// Creates a builder with no schemes registered.
114    ///
115    /// Useful when the caller wants explicit, scheme-by-scheme control over the compressor.
116    pub fn empty() -> Self {
117        Self {
118            schemes: Vec::new(),
119        }
120    }
121
122    /// Adds an external compression scheme not in [`ALL_SCHEMES`].
123    ///
124    /// This allows encoding crates outside of `vortex-btrblocks` to register their own schemes
125    /// with the compressor.
126    ///
127    /// # Panics
128    ///
129    /// Panics if a scheme with the same [`SchemeId`] is already present.
130    pub fn with_new_scheme(mut self, scheme: &'static dyn Scheme) -> Self {
131        assert!(
132            !self.schemes.iter().any(|s| s.id() == scheme.id()),
133            "scheme {:?} is already present in the builder",
134            scheme.id(),
135        );
136
137        self.schemes.push(scheme);
138        self
139    }
140
141    /// Adds compact encoding schemes (Zstd for strings and binary, Pco for numerics).
142    ///
143    /// This provides better compression ratios than the default, especially for floating-point
144    /// heavy datasets. Requires the `zstd` feature. When the `pco` feature is also enabled,
145    /// Pco schemes for integers and floats are included.
146    ///
147    /// # Panics
148    ///
149    /// Panics if any of the compact schemes are already present.
150    #[cfg(feature = "zstd")]
151    pub fn with_compact(self) -> Self {
152        let builder = self
153            .with_new_scheme(&string::ZstdScheme)
154            .with_new_scheme(&binary::ZstdScheme);
155
156        #[cfg(feature = "pco")]
157        let builder = builder
158            .with_new_scheme(&integer::PcoScheme)
159            .with_new_scheme(&float::PcoScheme);
160
161        builder
162    }
163
164    /// Excludes schemes without CUDA kernel support and adds Zstd for string and binary compression.
165    ///
166    /// With the `unstable_encodings` feature, buffer-level Zstd compression is used which
167    /// preserves the array buffer layout for zero-conversion GPU decompression. Without it,
168    /// interleaved Zstd compression is used.
169    ///
170    /// This preset is intended for files that will be decoded by CUDA kernels. It may choose a
171    /// larger encoded representation than the default compressor.
172    pub fn only_cuda_compatible(self) -> Self {
173        // String fragmentation schemes (OnPair, FSST) require host-side
174        // dictionary expansion at decode time, which is incompatible with
175        // pure-GPU decompression paths. Strip whichever string-fragment
176        // scheme is enabled by feature.
177        #[cfg_attr(not(feature = "unstable_encodings"), allow(unused_mut))]
178        let mut excluded: Vec<SchemeId> = vec![
179            integer::SparseScheme.id(),
180            integer::IntRLEScheme.id(),
181            float::FloatRLEScheme.id(),
182            float::NullDominatedSparseScheme.id(),
183            string::StringDictScheme.id(),
184            string::FSSTScheme.id(),
185            binary::BinaryDictScheme.id(),
186        ];
187        #[cfg(feature = "unstable_encodings")]
188        excluded.push(string::OnPairScheme.id());
189        // Delta has no GPU decode kernel and its prefix-sum decode is inherently sequential, so it
190        // is incompatible with pure-GPU decompression paths.
191        #[cfg(feature = "unstable_encodings")]
192        excluded.push(integer::DeltaScheme::default().id());
193        let builder = self.exclude_schemes(excluded);
194
195        #[cfg(all(feature = "zstd", feature = "unstable_encodings"))]
196        let builder = builder
197            .with_new_scheme(&string::ZstdBuffersScheme)
198            .with_new_scheme(&binary::ZstdBuffersScheme);
199        #[cfg(all(feature = "zstd", not(feature = "unstable_encodings")))]
200        let builder = builder
201            .with_new_scheme(&string::ZstdScheme)
202            .with_new_scheme(&binary::ZstdScheme);
203
204        builder
205    }
206
207    /// Removes the specified compression schemes by their [`SchemeId`].
208    pub fn exclude_schemes(mut self, ids: impl IntoIterator<Item = SchemeId>) -> Self {
209        let ids: HashSet<_> = ids.into_iter().collect();
210        self.schemes.retain(|s| !ids.contains(&s.id()));
211        self
212    }
213
214    /// Builds the configured [`BtrBlocksCompressor`].
215    pub fn build(self) -> BtrBlocksCompressor {
216        BtrBlocksCompressor(CascadingCompressor::new(self.schemes))
217    }
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223
224    #[test]
225    fn empty_starts_with_no_schemes() {
226        let builder = BtrBlocksCompressorBuilder::empty();
227        assert!(builder.schemes.is_empty());
228    }
229
230    #[test]
231    fn default_includes_all_schemes() {
232        let builder = BtrBlocksCompressorBuilder::default();
233        assert_eq!(builder.schemes.len(), ALL_SCHEMES.len());
234    }
235}
vortex_btrblocks/builder.rs

vortex_btrblocks/
builder.rs
