Skip to main content

krabka_compression/
codec_type.rs

1//! `CompressionType` enum mapping to Kafka's record-batch attribute bits.
2
3use crate::CompressionError;
4
5/// Kafka's gzip levels: `Deflater.BEST_SPEED`, `Deflater.BEST_COMPRESSION`
6/// and `Deflater.DEFAULT_COMPRESSION`.
7const GZIP_LEVELS: LevelRange = LevelRange {
8    min: 1,
9    max: 9,
10    default: -1,
11};
12
13/// Kafka's lz4 levels, from `net.jpountz.lz4.LZ4Constants`.
14const LZ4_LEVELS: LevelRange = LevelRange {
15    min: 1,
16    max: 17,
17    default: 9,
18};
19
20/// Kafka's zstd levels: `ZSTD_minCLevel`, `ZSTD_MAX_CLEVEL` and
21/// `ZSTD_CLEVEL_DEFAULT`.
22const ZSTD_LEVELS: LevelRange = LevelRange {
23    min: -131_072,
24    max: 22,
25    default: 3,
26};
27
28/// The levels of one codec.
29#[derive(Debug, Clone, Copy)]
30struct LevelRange {
31    min: i32,
32    max: i32,
33    default: i32,
34}
35
36/// Codec identifier matching the lowest three bits of Kafka's record-batch
37/// attribute byte.
38#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
39#[repr(u8)]
40#[non_exhaustive]
41pub enum CompressionType {
42    None = 0,
43    Gzip = 1,
44    Snappy = 2,
45    Lz4 = 3,
46    Zstd = 4,
47}
48
49impl CompressionType {
50    /// Decode the lowest three bits of a Kafka record-batch attribute byte.
51    ///
52    /// The method returns `None` for codec ids outside `0..=4`.
53    #[must_use]
54    pub fn from_attribute_bits(b: u8) -> Option<Self> {
55        match b & 0b0000_0111 {
56            0 => Some(Self::None),
57            1 => Some(Self::Gzip),
58            2 => Some(Self::Snappy),
59            3 => Some(Self::Lz4),
60            4 => Some(Self::Zstd),
61            _ => None,
62        }
63    }
64
65    /// Encode this codec into the lowest three bits of an attribute byte.
66    #[must_use]
67    pub fn as_attribute_bits(self) -> u8 {
68        self as u8
69    }
70
71    /// The codec name that Kafka uses in `compression.type`.
72    #[must_use]
73    pub const fn name(self) -> &'static str {
74        match self {
75            Self::None => "none",
76            Self::Gzip => "gzip",
77            Self::Snappy => "snappy",
78            Self::Lz4 => "lz4",
79            Self::Zstd => "zstd",
80        }
81    }
82
83    /// Kafka's default compression level of the codec
84    /// (`CompressionType.defaultLevel`), or `None` when the codec has no
85    /// levels.
86    #[must_use]
87    pub const fn default_level(self) -> Option<i32> {
88        match self.levels() {
89            Some(levels) => Some(levels.default),
90            None => None,
91        }
92    }
93
94    /// Kafka's lowest compression level of the codec
95    /// (`CompressionType.minLevel`), or `None` when the codec has no levels.
96    #[must_use]
97    pub const fn min_level(self) -> Option<i32> {
98        match self.levels() {
99            Some(levels) => Some(levels.min),
100            None => None,
101        }
102    }
103
104    /// Kafka's highest compression level of the codec
105    /// (`CompressionType.maxLevel`), or `None` when the codec has no levels.
106    #[must_use]
107    pub const fn max_level(self) -> Option<i32> {
108        match self.levels() {
109            Some(levels) => Some(levels.max),
110            None => None,
111        }
112    }
113
114    /// Check `level` as Kafka's `CompressionType.levelValidator` does (KIP-390).
115    ///
116    /// Gzip accepts 1 to 9, and also -1, which is the zlib default. Lz4 and
117    /// zstd accept their range only.
118    ///
119    /// # Errors
120    ///
121    /// Returns [`CompressionError::InvalidLevel`] with Kafka's message when the
122    /// level is outside the range, or when the codec has no levels.
123    pub fn check_level(self, level: i32) -> Result<(), CompressionError> {
124        let invalid = |reason: String| CompressionError::InvalidLevel {
125            codec: self.name(),
126            level,
127            reason,
128        };
129        match (self, self.levels()) {
130            (_, None) => Err(invalid(format!(
131                "Compression levels are not defined for this compression type: {}",
132                self.name()
133            ))),
134            (Self::Gzip, Some(levels)) => {
135                if level > levels.max || (level < levels.min && level != levels.default) {
136                    Err(invalid(format!(
137                        "Value must be between {} and {} or equal to {}",
138                        levels.min, levels.max, levels.default
139                    )))
140                } else {
141                    Ok(())
142                }
143            }
144            (_, Some(levels)) => {
145                if level < levels.min {
146                    Err(invalid(format!("Value must be at least {}", levels.min)))
147                } else if level > levels.max {
148                    Err(invalid(format!(
149                        "Value must be no more than {}",
150                        levels.max
151                    )))
152                } else {
153                    Ok(())
154                }
155            }
156        }
157    }
158
159    const fn levels(self) -> Option<LevelRange> {
160        match self {
161            Self::None | Self::Snappy => None,
162            Self::Gzip => Some(GZIP_LEVELS),
163            Self::Lz4 => Some(LZ4_LEVELS),
164            Self::Zstd => Some(ZSTD_LEVELS),
165        }
166    }
167}
168
169#[cfg(test)]
170mod tests {
171
172    use super::*;
173
174    #[test]
175    fn attribute_bits_roundtrip() {
176        for (_name, compression) in [
177            ("none", CompressionType::None),
178            ("gzip", CompressionType::Gzip),
179            ("snappy", CompressionType::Snappy),
180            ("lz4", CompressionType::Lz4),
181            ("zstd", CompressionType::Zstd),
182        ] {
183            assert2::assert!(
184                CompressionType::from_attribute_bits(compression.as_attribute_bits())
185                    == Some(compression)
186            );
187        }
188    }
189
190    /// Each row follows Kafka's `CompressionType` level constants and
191    /// `levelValidator` (gzip), or `ConfigDef.Range.ensureValid` (lz4, zstd).
192    #[test]
193    fn levels_follow_kafka_compression_type() {
194        let reasons = |codec: CompressionType, levels: &[i32]| -> Vec<Option<String>> {
195            levels
196                .iter()
197                .map(|&level| match codec.check_level(level) {
198                    Ok(()) => None,
199                    Err(CompressionError::InvalidLevel { reason, .. }) => Some(reason),
200                    Err(other) => Some(format!("unexpected error {other}")),
201                })
202                .collect()
203        };
204        let gzip_reason = Some("Value must be between 1 and 9 or equal to -1".to_owned());
205        let none_levels = |name: &str| {
206            Some(format!(
207                "Compression levels are not defined for this compression type: {name}"
208            ))
209        };
210        let actual = [
211            (
212                "none",
213                CompressionType::None.default_level(),
214                CompressionType::None.min_level(),
215                CompressionType::None.max_level(),
216                reasons(CompressionType::None, &[0]),
217            ),
218            (
219                "gzip",
220                CompressionType::Gzip.default_level(),
221                CompressionType::Gzip.min_level(),
222                CompressionType::Gzip.max_level(),
223                reasons(CompressionType::Gzip, &[-2, -1, 0, 1, 9, 10]),
224            ),
225            (
226                "snappy",
227                CompressionType::Snappy.default_level(),
228                CompressionType::Snappy.min_level(),
229                CompressionType::Snappy.max_level(),
230                reasons(CompressionType::Snappy, &[1]),
231            ),
232            (
233                "lz4",
234                CompressionType::Lz4.default_level(),
235                CompressionType::Lz4.min_level(),
236                CompressionType::Lz4.max_level(),
237                reasons(CompressionType::Lz4, &[0, 1, 17, 18]),
238            ),
239            (
240                "zstd",
241                CompressionType::Zstd.default_level(),
242                CompressionType::Zstd.min_level(),
243                CompressionType::Zstd.max_level(),
244                reasons(CompressionType::Zstd, &[-131_073, -131_072, 22, 23]),
245            ),
246        ];
247        let expected = [
248            ("none", None, None, None, vec![none_levels("none")]),
249            (
250                "gzip",
251                Some(-1),
252                Some(1),
253                Some(9),
254                vec![
255                    gzip_reason.clone(),
256                    None,
257                    gzip_reason.clone(),
258                    None,
259                    None,
260                    gzip_reason,
261                ],
262            ),
263            ("snappy", None, None, None, vec![none_levels("snappy")]),
264            (
265                "lz4",
266                Some(9),
267                Some(1),
268                Some(17),
269                vec![
270                    Some("Value must be at least 1".to_owned()),
271                    None,
272                    None,
273                    Some("Value must be no more than 17".to_owned()),
274                ],
275            ),
276            (
277                "zstd",
278                Some(3),
279                Some(-131_072),
280                Some(22),
281                vec![
282                    Some("Value must be at least -131072".to_owned()),
283                    None,
284                    None,
285                    Some("Value must be no more than 22".to_owned()),
286                ],
287            ),
288        ];
289        assert2::assert!(actual == expected);
290    }
291
292    #[test]
293    fn attribute_bits_mask() {
294        // Only the low 3 bits define the codec; upper bits are other flags.
295        assert2::assert!(
296            CompressionType::from_attribute_bits(0b1111_1000 | 0b0000_0001)
297                == Some(CompressionType::Gzip)
298        );
299    }
300
301    #[test]
302    fn attribute_bits_unknown() {
303        for (_name, bits) in [("reserved five", 5), ("reserved seven", 7)] {
304            assert2::assert!(CompressionType::from_attribute_bits(bits) == None);
305        }
306    }
307}