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 itertools::Itertools;
7use vortex_utils::aliases::hash_set::HashSet;
8
9use crate::BtrBlocksCompressor;
10use crate::FloatCode;
11use crate::IntCode;
12use crate::StringCode;
13use crate::compressor::float::ALL_FLOAT_SCHEMES;
14use crate::compressor::float::FloatScheme;
15use crate::compressor::integer::ALL_INT_SCHEMES;
16use crate::compressor::integer::IntegerScheme;
17use crate::compressor::string::ALL_STRING_SCHEMES;
18use crate::compressor::string::StringScheme;
19
20/// Builder for creating configured [`BtrBlocksCompressor`] instances.
21///
22/// Use this builder to configure which compression schemes are allowed for each data type.
23/// By default, all schemes are enabled.
24///
25/// # Examples
26///
27/// ```rust
28/// use vortex_btrblocks::{BtrBlocksCompressorBuilder, IntCode, FloatCode};
29///
30/// // Default compressor - all schemes allowed
31/// let compressor = BtrBlocksCompressorBuilder::default().build();
32///
33/// // Exclude specific schemes
34/// let compressor = BtrBlocksCompressorBuilder::default()
35///     .exclude_int([IntCode::Dict])
36///     .build();
37///
38/// // Exclude then re-include
39/// let compressor = BtrBlocksCompressorBuilder::default()
40///     .exclude_int([IntCode::Dict, IntCode::Rle])
41///     .include_int([IntCode::Dict])
42///     .build();
43/// ```
44#[derive(Debug, Clone)]
45pub struct BtrBlocksCompressorBuilder {
46    int_schemes: HashSet<&'static dyn IntegerScheme>,
47    float_schemes: HashSet<&'static dyn FloatScheme>,
48    string_schemes: HashSet<&'static dyn StringScheme>,
49}
50
51impl Default for BtrBlocksCompressorBuilder {
52    fn default() -> Self {
53        Self {
54            int_schemes: ALL_INT_SCHEMES
55                .iter()
56                .copied()
57                .filter(|s| s.code() != IntCode::Pco)
58                .collect(),
59            float_schemes: ALL_FLOAT_SCHEMES
60                .iter()
61                .copied()
62                .filter(|s| s.code() != FloatCode::Pco)
63                .collect(),
64            string_schemes: ALL_STRING_SCHEMES
65                .iter()
66                .copied()
67                .filter(|s| s.code() != StringCode::Zstd && s.code() != StringCode::ZstdBuffers)
68                .collect(),
69        }
70    }
71}
72
73impl BtrBlocksCompressorBuilder {
74    /// Create a new builder with no encodings enabled.
75    pub fn empty() -> Self {
76        Self {
77            int_schemes: Default::default(),
78            float_schemes: Default::default(),
79            string_schemes: Default::default(),
80        }
81    }
82
83    /// Excludes the specified integer compression schemes.
84    pub fn exclude_int(mut self, codes: impl IntoIterator<Item = IntCode>) -> Self {
85        let codes: HashSet<_> = codes.into_iter().collect();
86        self.int_schemes.retain(|s| !codes.contains(&s.code()));
87        self
88    }
89
90    /// Excludes the specified float compression schemes.
91    pub fn exclude_float(mut self, codes: impl IntoIterator<Item = FloatCode>) -> Self {
92        let codes: HashSet<_> = codes.into_iter().collect();
93        self.float_schemes.retain(|s| !codes.contains(&s.code()));
94        self
95    }
96
97    /// Excludes the specified string compression schemes.
98    pub fn exclude_string(mut self, codes: impl IntoIterator<Item = StringCode>) -> Self {
99        let codes: HashSet<_> = codes.into_iter().collect();
100        self.string_schemes.retain(|s| !codes.contains(&s.code()));
101        self
102    }
103
104    /// Includes the specified integer compression schemes.
105    pub fn include_int(mut self, codes: impl IntoIterator<Item = IntCode>) -> Self {
106        let codes: HashSet<_> = codes.into_iter().collect();
107        for scheme in ALL_INT_SCHEMES {
108            if codes.contains(&scheme.code()) {
109                self.int_schemes.insert(*scheme);
110            }
111        }
112        self
113    }
114
115    /// Includes the specified float compression schemes.
116    pub fn include_float(mut self, codes: impl IntoIterator<Item = FloatCode>) -> Self {
117        let codes: HashSet<_> = codes.into_iter().collect();
118        for scheme in ALL_FLOAT_SCHEMES {
119            if codes.contains(&scheme.code()) {
120                self.float_schemes.insert(*scheme);
121            }
122        }
123        self
124    }
125
126    /// Includes the specified string compression schemes.
127    pub fn include_string(mut self, codes: impl IntoIterator<Item = StringCode>) -> Self {
128        let codes: HashSet<_> = codes.into_iter().collect();
129        for scheme in ALL_STRING_SCHEMES {
130            if codes.contains(&scheme.code()) {
131                self.string_schemes.insert(*scheme);
132            }
133        }
134        self
135    }
136
137    /// Builds the configured `BtrBlocksCompressor`.
138    pub fn build(self) -> BtrBlocksCompressor {
139        // Note we should apply the schemes in the same order, in case try conflict.
140        BtrBlocksCompressor {
141            int_schemes: self
142                .int_schemes
143                .into_iter()
144                .sorted_by_key(|s| s.code())
145                .collect_vec(),
146            float_schemes: self
147                .float_schemes
148                .into_iter()
149                .sorted_by_key(|s| s.code())
150                .collect_vec(),
151            string_schemes: self
152                .string_schemes
153                .into_iter()
154                .sorted_by_key(|s| s.code())
155                .collect_vec(),
156        }
157    }
158}
vortex_btrblocks/builder.rs

vortex_btrblocks/
builder.rs
