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

vortex_btrblocks/
builder.rs
