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_decimal_byte_parts::decimal_byte_parts_v2_id;
8use vortex_fastlanes::for_v2_id;
9use vortex_session::VortexSession;
10use vortex_utils::aliases::hash_set::HashSet;
11
12use crate::BtrBlocksCompressor;
13use crate::CascadingCompressor;
14use crate::CompressionSessionExt;
15use crate::Scheme;
16use crate::SchemeExt;
17use crate::SchemeId;
18use crate::allowed_ids::AllowedIds;
19use crate::schemes::binary;
20use crate::schemes::float;
21use crate::schemes::integer;
22use crate::schemes::string;
23
24/// The preset a [`BtrBlocksCompressorBuilder`] builds with.
25///
26/// Every mode except [`All`](Self::All) excludes some schemes in
27/// [`build`](BtrBlocksCompressorBuilder::build).
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub(crate) enum CompressionMode {
30    /// Excludes nothing. Set by [`empty`](BtrBlocksCompressorBuilder::empty).
31    All,
32    /// Excludes Zstd and Pco. Set by [`from_session`](BtrBlocksCompressorBuilder::from_session).
33    Default,
34    /// Excludes buffer-level Zstd, keeping Zstd for strings and binary and Pco for numerics.
35    /// Set by [`with_compact`](BtrBlocksCompressorBuilder::with_compact).
36    Compact,
37    /// Excludes schemes without CUDA kernel support, keeping FSST for strings and both Zstd
38    /// schemes for binary. Set by
39    /// [`only_cuda_compatible`](BtrBlocksCompressorBuilder::only_cuda_compatible).
40    Cuda,
41}
42
43impl CompressionMode {
44    /// Returns the schemes [`build`](BtrBlocksCompressorBuilder::build) drops in this mode.
45    fn excluded_schemes(self) -> Vec<SchemeId> {
46        let mut excluded = Vec::new();
47        match self {
48            Self::All => {}
49            Self::Default => {
50                #[cfg(feature = "zstd")]
51                excluded.extend([
52                    string::ZstdScheme.id(),
53                    binary::ZstdScheme.id(),
54                    binary::ZstdBuffersScheme.id(),
55                ]);
56                #[cfg(feature = "pco")]
57                excluded.extend([integer::PcoScheme.id(), float::PcoScheme.id()]);
58            }
59            Self::Compact => {
60                #[cfg(feature = "zstd")]
61                excluded.push(binary::ZstdBuffersScheme.id());
62            }
63            Self::Cuda => {
64                // Keep FSST, which has a CUDA decoder and direct Arrow offset-based export. Other
65                // string fragmentation and dictionary schemes still require unsupported decode
66                // paths. Delta has a CUDA decode kernel, but stays excluded until GPU delta
67                // decode is benchmarked against the schemes it would displace.
68                excluded.extend([
69                    integer::DeltaScheme::default().id(),
70                    integer::SparseScheme.id(),
71                    integer::IntRLEScheme.id(),
72                    float::ALPRDScheme.id(),
73                    float::FloatRLEScheme.id(),
74                    float::NullDominatedSparseScheme.id(),
75                    string::NullDominatedSparseScheme.id(),
76                    string::StringDictScheme.id(),
77                    binary::BinaryDictScheme.id(),
78                ]);
79                // Both binary Zstd schemes are kept. Buffer-level compression preserves binary
80                // arrays' buffer layout for zero-conversion GPU decompression, but belongs to the
81                // opt-in `zstd` edition, so the session's enabled editions decide which of the
82                // two is kept.
83                #[cfg(feature = "zstd")]
84                excluded.push(string::ZstdScheme.id());
85                #[cfg(feature = "pco")]
86                excluded.extend([integer::PcoScheme.id(), float::PcoScheme.id()]);
87            }
88        }
89        excluded
90    }
91
92    /// Returns the serialized IDs [`build`](BtrBlocksCompressorBuilder::build) disallows in this
93    /// mode, on top of the session's restrictions.
94    fn excluded_encodings(self) -> Vec<ArrayId> {
95        match self {
96            Self::All | Self::Default | Self::Compact => Vec::new(),
97            // Multi-part DecimalByteParts arrays and FoR arrays with per-chunk references have no
98            // CUDA decode kernel.
99            Self::Cuda => vec![decimal_byte_parts_v2_id(), for_v2_id()],
100        }
101    }
102}
103
104/// Builder for creating configured [`BtrBlocksCompressor`] instances.
105///
106/// [`from_session`](Self::from_session) starts from the schemes registered in the session's
107/// [`CompressionSession`](crate::CompressionSession), in registration order, and
108/// [`build`](Self::build) excludes some of them: by default Zstd and Pco.
109/// [`with_compact`](Self::with_compact) and [`only_cuda_compatible`](Self::only_cuda_compatible)
110/// switch to a different preset.
111///
112/// The builder also tracks which serialized array IDs its schemes may produce, taken from the
113/// session's registered arrays and enabled editions. [`build`](Self::build) drops every scheme
114/// that produces a disallowed ID.
115///
116/// # Examples
117///
118/// ```rust
119/// use vortex_btrblocks::{BtrBlocksCompressorBuilder, CompressionSession, Scheme, SchemeExt};
120/// use vortex_btrblocks::schemes::integer::IntDictScheme;
121/// use vortex_session::VortexSession;
122///
123/// let session = VortexSession::empty().with::<CompressionSession>();
124///
125/// // Compressor with the session's schemes, restricted to the encodings it allows. This session
126/// // enables no editions, so every scheme is dropped.
127/// let compressor = BtrBlocksCompressorBuilder::from_session(&session).build();
128///
129/// // Remove specific schemes.
130/// let compressor = BtrBlocksCompressorBuilder::from_session(&session)
131///     .exclude_schemes([IntDictScheme.id()])
132///     .build();
133/// ```
134#[derive(Debug, Clone)]
135pub struct BtrBlocksCompressorBuilder {
136    schemes: Vec<&'static dyn Scheme>,
137    allowed: AllowedIds,
138    mode: CompressionMode,
139}
140
141impl BtrBlocksCompressorBuilder {
142    /// Creates a builder with every scheme registered in the session's
143    /// [`CompressionSession`](crate::CompressionSession), allowing the serialized IDs registered
144    /// in the session and permitted by its enabled editions.
145    pub fn from_session(session: &VortexSession) -> Self {
146        Self {
147            schemes: session.compression().schemes().to_vec(),
148            allowed: AllowedIds::from_session(session),
149            mode: CompressionMode::Default,
150        }
151    }
152
153    /// Creates a builder with no schemes registered.
154    ///
155    /// Useful when the caller wants explicit, scheme-by-scheme control over the compressor.
156    /// Every added scheme and every serialized ID is allowed.
157    pub fn empty() -> Self {
158        Self {
159            schemes: Vec::new(),
160            allowed: AllowedIds::all(),
161            mode: CompressionMode::All,
162        }
163    }
164
165    /// Adds a compression scheme not registered on the session.
166    ///
167    /// This allows encoding crates outside of `vortex-btrblocks` to register their own schemes
168    /// with the compressor.
169    ///
170    /// # Panics
171    ///
172    /// Panics if a scheme with the same [`SchemeId`] is already present.
173    pub fn with_new_scheme(mut self, scheme: &'static dyn Scheme) -> Self {
174        assert!(
175            !self.schemes.iter().any(|s| s.id() == scheme.id()),
176            "scheme {:?} is already present in the builder",
177            scheme.id(),
178        );
179
180        self.schemes.push(scheme);
181        self
182    }
183
184    /// Switches to the compact preset, keeping Zstd for strings and binary and Pco for numerics.
185    ///
186    /// This provides better compression ratios than the default, especially for floating-point
187    /// heavy datasets. The Zstd and Pco schemes are only registered with the `zstd` and `pco`
188    /// features.
189    pub fn with_compact(mut self) -> Self {
190        self.mode = CompressionMode::Compact;
191        self
192    }
193
194    /// Switches to the CUDA preset, excluding schemes without CUDA kernel support, keeping FSST
195    /// for string compression and Zstd for binary compression.
196    ///
197    /// This preset is intended for files that will be decoded by CUDA kernels. It may choose a
198    /// larger encoded representation than the default compressor.
199    pub fn only_cuda_compatible(mut self) -> Self {
200        self.mode = CompressionMode::Cuda;
201        self
202    }
203
204    /// Removes the specified compression schemes by their [`SchemeId`].
205    pub fn exclude_schemes(mut self, ids: impl IntoIterator<Item = SchemeId>) -> Self {
206        let ids: HashSet<_> = ids.into_iter().collect();
207        self.schemes.retain(|s| !ids.contains(&s.id()));
208        self
209    }
210
211    /// Allows serialized IDs the session's enabled editions do not permit, still requiring them
212    /// to be registered in the session.
213    pub fn disable_editions(mut self) -> Self {
214        self.allowed.editions = None;
215        self
216    }
217
218    /// Allows serialized IDs regardless of the session's registered arrays and enabled editions.
219    ///
220    /// Serialized IDs excluded by a preset stay excluded.
221    pub fn unrestricted(mut self) -> Self {
222        self.allowed.registered = None;
223        self.allowed.editions = None;
224        self
225    }
226
227    /// Builds the configured [`BtrBlocksCompressor`] from the schemes that its preset does not
228    /// exclude and whose produced serialized IDs are all allowed.
229    pub fn build(self) -> BtrBlocksCompressor {
230        BtrBlocksCompressor(CascadingCompressor::new(self.allowed_schemes()))
231    }
232
233    fn allowed_schemes(&self) -> Vec<&'static dyn Scheme> {
234        let excluded: HashSet<SchemeId> = self.mode.excluded_schemes().into_iter().collect();
235        let excluded_ids: HashSet<ArrayId> = self.mode.excluded_encodings().into_iter().collect();
236        let allowed = |id: &ArrayId| self.allowed.is_allowed(id) && !excluded_ids.contains(id);
237        self.schemes
238            .iter()
239            .copied()
240            .filter(|s| !excluded.contains(&s.id()))
241            .map(|s| s.refine(&allowed))
242            .filter(|s| s.produced_encodings().iter().all(&allowed))
243            .collect()
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use rstest::rstest;
250    use vortex_array::ArrayId;
251    use vortex_array::VTable;
252    use vortex_array::arrays::VarBin;
253    use vortex_decimal_byte_parts::DecimalByteParts;
254    use vortex_decimal_byte_parts::decimal_byte_parts_v1_id;
255    use vortex_fsst::FSST;
256
257    use super::*;
258    use crate::CompressionSession;
259    use crate::schemes::decimal::DECIMAL_V1;
260    use crate::schemes::decimal::DECIMAL_V2;
261
262    #[rstest]
263    #[case::fsst_missing_codes(&string::FSSTScheme, vec![FSST.id()], false)]
264    #[case::fsst_with_codes(&string::FSSTScheme, vec![FSST.id(), VarBin.id()], true)]
265    #[case::decimal_v1_id(&DECIMAL_V1, vec![decimal_byte_parts_v1_id()], true)]
266    #[case::decimal_both_ids(&DECIMAL_V2, vec![decimal_byte_parts_v1_id(), DecimalByteParts.id()], true)]
267    fn test_allowed_schemes(
268        #[case] scheme: &'static dyn Scheme,
269        #[case] allowed: Vec<ArrayId>,
270        #[case] retained: bool,
271    ) {
272        let mut builder = BtrBlocksCompressorBuilder::empty().with_new_scheme(scheme);
273        builder.allowed.editions = Some(allowed.into_iter().collect());
274
275        assert_eq!(!builder.allowed_schemes().is_empty(), retained);
276    }
277
278    #[test]
279    fn delta_is_excluded_only_by_cuda() {
280        let session = VortexSession::empty().with::<CompressionSession>();
281        let builder = BtrBlocksCompressorBuilder::from_session(&session).unrestricted();
282        let has_delta = |builder: BtrBlocksCompressorBuilder| {
283            builder
284                .allowed_schemes()
285                .iter()
286                .any(|s| s.id() == integer::DeltaScheme::default().id())
287        };
288        assert!(has_delta(builder.clone()));
289        assert!(has_delta(builder.clone().with_compact()));
290        assert!(!has_delta(builder.only_cuda_compatible()));
291    }
292}
vortex_btrblocks/builder.rs

vortex_btrblocks/
builder.rs
