iceberg/
compression.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Compression codec support for data compression and decompression.
19
20use std::fmt;
21use std::io::{Read, Write};
22
23use flate2::Compression;
24use flate2::read::GzDecoder;
25use flate2::write::GzEncoder;
26use serde::{Deserialize, Deserializer, Serialize, Serializer};
27
28use crate::{Error, ErrorKind, Result};
29
30/// Default compression level for Zstandard (zstd).
31const ZSTD_DEFAULT_LEVEL: u8 = 3;
32/// Default compression level for Gzip.
33const GZIP_DEFAULT_LEVEL: u8 = 6;
34/// Maximum compression level for Gzip.
35const GZIP_MAX_LEVEL: u8 = 9;
36/// Default compression level for Brotli.
37const BROTLI_DEFAULT_LEVEL: u8 = 1;
38
39/// Data compression formats
40#[derive(Debug, PartialEq, Eq, Clone, Copy, Default)]
41pub enum CompressionCodec {
42    #[default]
43    /// No compression
44    None,
45    /// LZ4 single compression frame with content size present
46    Lz4,
47    /// LZ4 raw block compression (no frame). Used by Parquet as `lz4_raw`.
48    Lz4Raw,
49    /// Zstandard single compression frame with content size present.
50    /// Level range is 0–22, where 0 means default compression level (not no compression).
51    /// Use [`CompressionCodec::zstd_default`] to construct with the default level.
52    Zstd(u8),
53    /// Gzip compression. Level range is 0–9, where 0 means no compression.
54    /// Use [`CompressionCodec::gzip_default`] to construct with the default level.
55    Gzip(u8),
56    /// Brotli compression. Level range is 0–11.
57    /// Use [`CompressionCodec::brotli_default`] to construct with the default level.
58    Brotli(u8),
59    /// LZO compression
60    Lzo,
61    /// Snappy compression
62    Snappy,
63}
64
65impl CompressionCodec {
66    /// Returns a Zstd codec with the default compression level.
67    pub const fn zstd_default() -> Self {
68        CompressionCodec::Zstd(ZSTD_DEFAULT_LEVEL)
69    }
70
71    /// Returns a Gzip codec with the default compression level.
72    pub const fn gzip_default() -> Self {
73        CompressionCodec::Gzip(GZIP_DEFAULT_LEVEL)
74    }
75
76    /// Returns a Brotli codec with the default compression level.
77    pub const fn brotli_default() -> Self {
78        CompressionCodec::Brotli(BROTLI_DEFAULT_LEVEL)
79    }
80
81    /// Returns the codec name as used in serialization and error messages.
82    pub fn name(&self) -> &'static str {
83        match self {
84            CompressionCodec::None => "none",
85            CompressionCodec::Lz4 => "lz4",
86            CompressionCodec::Lz4Raw => "lz4_raw",
87            CompressionCodec::Zstd(_) => "zstd",
88            CompressionCodec::Gzip(_) => "gzip",
89            CompressionCodec::Brotli(_) => "brotli",
90            CompressionCodec::Lzo => "lzo",
91            CompressionCodec::Snappy => "snappy",
92        }
93    }
94}
95
96// Note: serialize/deserialize do not round-trip the compression level. Iceberg configuration
97// only the codec name (e.g. "zstd"), not the level, so deserialization always produces the
98// default level. A `Zstd(5)` written to metadata will be read back as `Zstd(3)`. Some
99// compression configuration (e.g. Avro metadata) has a separate level field alongside the codec name.
100impl Serialize for CompressionCodec {
101    fn serialize<S: Serializer>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> {
102        serializer.serialize_str(self.name())
103    }
104}
105
106impl<'de> Deserialize<'de> for CompressionCodec {
107    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> std::result::Result<Self, D::Error> {
108        let s = String::deserialize(deserializer)?;
109        match s.to_lowercase().as_str() {
110            "none" | "uncompressed" => Ok(CompressionCodec::None),
111            "lz4" => Ok(CompressionCodec::Lz4),
112            "lz4_raw" => Ok(CompressionCodec::Lz4Raw),
113            "zstd" => Ok(CompressionCodec::zstd_default()),
114            "gzip" => Ok(CompressionCodec::gzip_default()),
115            "brotli" => Ok(CompressionCodec::brotli_default()),
116            "lzo" => Ok(CompressionCodec::Lzo),
117            "snappy" => Ok(CompressionCodec::Snappy),
118            other => Err(serde::de::Error::unknown_variant(other, &[
119                "none",
120                "uncompressed",
121                "lz4",
122                "lz4_raw",
123                "zstd",
124                "gzip",
125                "brotli",
126                "lzo",
127                "snappy",
128            ])),
129        }
130    }
131}
132
133impl fmt::Display for CompressionCodec {
134    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
135        match self {
136            CompressionCodec::None => write!(f, "None"),
137            CompressionCodec::Lz4 => write!(f, "Lz4"),
138            CompressionCodec::Lz4Raw => write!(f, "Lz4Raw"),
139            CompressionCodec::Zstd(level) => write!(f, "Zstd(level={level})"),
140            CompressionCodec::Gzip(level) => write!(f, "Gzip(level={level})"),
141            CompressionCodec::Brotli(level) => write!(f, "Brotli(level={level})"),
142            CompressionCodec::Lzo => write!(f, "Lzo"),
143            CompressionCodec::Snappy => write!(f, "Snappy"),
144        }
145    }
146}
147
148impl CompressionCodec {
149    pub(crate) fn decompress(&self, bytes: Vec<u8>) -> Result<Vec<u8>> {
150        match self {
151            CompressionCodec::None => Ok(bytes),
152            CompressionCodec::Lz4 => Err(Error::new(
153                ErrorKind::FeatureUnsupported,
154                "LZ4 decompression is not supported currently",
155            )),
156            CompressionCodec::Zstd(_) => Ok(zstd::stream::decode_all(&bytes[..])?),
157            CompressionCodec::Gzip(_) => {
158                let mut decoder = GzDecoder::new(&bytes[..]);
159                let mut decompressed = Vec::new();
160                decoder.read_to_end(&mut decompressed)?;
161                Ok(decompressed)
162            }
163            CompressionCodec::Lz4Raw => Err(Error::new(
164                ErrorKind::FeatureUnsupported,
165                "LZ4_RAW decompression is not supported currently",
166            )),
167            CompressionCodec::Brotli(_) => Err(Error::new(
168                ErrorKind::FeatureUnsupported,
169                "Brotli decompression is not supported currently",
170            )),
171            CompressionCodec::Lzo => Err(Error::new(
172                ErrorKind::FeatureUnsupported,
173                "LZO decompression is not supported currently",
174            )),
175            CompressionCodec::Snappy => Err(Error::new(
176                ErrorKind::FeatureUnsupported,
177                "Snappy decompression is not supported currently",
178            )),
179        }
180    }
181
182    pub(crate) fn compress(&self, bytes: Vec<u8>) -> Result<Vec<u8>> {
183        match self {
184            CompressionCodec::None => Ok(bytes),
185            CompressionCodec::Lz4 => Err(Error::new(
186                ErrorKind::FeatureUnsupported,
187                "LZ4 compression is not supported currently",
188            )),
189            CompressionCodec::Zstd(level) => {
190                let writer = Vec::<u8>::new();
191                let mut encoder = zstd::stream::Encoder::new(writer, *level as i32)?;
192                encoder.include_checksum(true)?;
193                encoder.set_pledged_src_size(Some(bytes.len().try_into()?))?;
194                std::io::copy(&mut &bytes[..], &mut encoder)?;
195                Ok(encoder.finish()?)
196            }
197            CompressionCodec::Gzip(level) => {
198                let compression = Compression::new((*level).min(GZIP_MAX_LEVEL) as u32);
199                let mut encoder = GzEncoder::new(Vec::new(), compression);
200                encoder.write_all(&bytes)?;
201                Ok(encoder.finish()?)
202            }
203            CompressionCodec::Lz4Raw => Err(Error::new(
204                ErrorKind::FeatureUnsupported,
205                "LZ4_RAW compression is not supported currently",
206            )),
207            CompressionCodec::Brotli(_) => Err(Error::new(
208                ErrorKind::FeatureUnsupported,
209                "Brotli compression is not supported currently",
210            )),
211            CompressionCodec::Lzo => Err(Error::new(
212                ErrorKind::FeatureUnsupported,
213                "LZO compression is not supported currently",
214            )),
215            CompressionCodec::Snappy => Err(Error::new(
216                ErrorKind::FeatureUnsupported,
217                "Snappy compression is not supported currently",
218            )),
219        }
220    }
221
222    pub(crate) fn is_none(&self) -> bool {
223        matches!(self, CompressionCodec::None)
224    }
225
226    /// Returns the file extension suffix for this compression codec.
227    /// Returns empty string for None, ".gz" for Gzip.
228    ///
229    /// # Errors
230    ///
231    /// Returns an error for Lz4 and Zstd as they are not fully supported.
232    pub fn suffix(&self) -> Result<&'static str> {
233        match self {
234            CompressionCodec::None => Ok(""),
235            CompressionCodec::Gzip(_) => Ok(".gz"),
236            codec @ (CompressionCodec::Lz4
237            | CompressionCodec::Lz4Raw
238            | CompressionCodec::Zstd(_)
239            | CompressionCodec::Brotli(_)
240            | CompressionCodec::Lzo
241            | CompressionCodec::Snappy) => Err(Error::new(
242                ErrorKind::FeatureUnsupported,
243                format!("suffix not defined for {codec:?}"),
244            )),
245        }
246    }
247}
248
249#[cfg(test)]
250mod tests {
251    use super::CompressionCodec;
252
253    #[tokio::test]
254    async fn test_compression_codec_none() {
255        let bytes_vec = [0_u8; 100].to_vec();
256
257        let codec = CompressionCodec::None;
258        let compressed = codec.compress(bytes_vec.clone()).unwrap();
259        assert_eq!(bytes_vec, compressed);
260        let decompressed = codec.decompress(compressed).unwrap();
261        assert_eq!(bytes_vec, decompressed);
262    }
263
264    #[tokio::test]
265    async fn test_compression_codec_compress() {
266        let bytes_vec = [0_u8; 100].to_vec();
267
268        let compression_codecs = [
269            CompressionCodec::zstd_default(),
270            CompressionCodec::gzip_default(),
271        ];
272
273        for codec in compression_codecs {
274            let compressed = codec.compress(bytes_vec.clone()).unwrap();
275            assert!(compressed.len() < bytes_vec.len());
276            let decompressed = codec.decompress(compressed).unwrap();
277            assert_eq!(decompressed, bytes_vec);
278        }
279    }
280
281    #[tokio::test]
282    async fn test_compression_codec_unsupported() {
283        let unsupported_codecs = [
284            (CompressionCodec::Lz4, "LZ4"),
285            (CompressionCodec::Snappy, "Snappy"),
286        ];
287        let bytes_vec = [0_u8; 100].to_vec();
288
289        for (codec, name) in unsupported_codecs {
290            assert_eq!(
291                codec.compress(bytes_vec.clone()).unwrap_err().to_string(),
292                format!("FeatureUnsupported => {name} compression is not supported currently"),
293            );
294
295            assert_eq!(
296                codec.decompress(bytes_vec.clone()).unwrap_err().to_string(),
297                format!("FeatureUnsupported => {name} decompression is not supported currently"),
298            );
299        }
300    }
301
302    #[test]
303    fn test_suffix() {
304        assert_eq!(CompressionCodec::None.suffix().unwrap(), "");
305        assert_eq!(CompressionCodec::gzip_default().suffix().unwrap(), ".gz");
306
307        assert!(CompressionCodec::Lz4.suffix().is_err());
308        assert!(CompressionCodec::zstd_default().suffix().is_err());
309        assert!(CompressionCodec::Snappy.suffix().is_err());
310
311        let lz4_err = CompressionCodec::Lz4.suffix().unwrap_err();
312        assert!(lz4_err.to_string().contains("suffix not defined for Lz4"));
313
314        let zstd_err = CompressionCodec::zstd_default().suffix().unwrap_err();
315        assert!(zstd_err.to_string().contains("suffix not defined for Zstd"));
316    }
317
318    #[test]
319    fn test_display() {
320        assert_eq!(CompressionCodec::None.to_string(), "None");
321        assert_eq!(CompressionCodec::Lz4.to_string(), "Lz4");
322        assert_eq!(
323            CompressionCodec::zstd_default().to_string(),
324            "Zstd(level=3)"
325        );
326        assert_eq!(CompressionCodec::Zstd(5).to_string(), "Zstd(level=5)");
327        assert_eq!(
328            CompressionCodec::gzip_default().to_string(),
329            "Gzip(level=6)"
330        );
331        assert_eq!(CompressionCodec::Gzip(9).to_string(), "Gzip(level=9)");
332        assert_eq!(CompressionCodec::Snappy.to_string(), "Snappy");
333    }
334}
iceberg/compression.rs

iceberg/
compression.rs
