vendor/iceberg-arrow59/src/compression.rs
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Compression codec support for data compression and decompression.
use std::fmt;
use std::io::{Read, Write};
use flate2::Compression;
use flate2::read::GzDecoder;
use flate2::write::GzEncoder;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use crate::{Error, ErrorKind, Result};
/// Default compression level for Zstandard (zstd).
const ZSTD_DEFAULT_LEVEL: u8 = 3;
/// Default compression level for Gzip.
const GZIP_DEFAULT_LEVEL: u8 = 6;
/// Maximum compression level for Gzip.
const GZIP_MAX_LEVEL: u8 = 9;
/// Default compression level for Brotli.
const BROTLI_DEFAULT_LEVEL: u8 = 1;
/// Data compression formats
#[derive(Debug, PartialEq, Eq, Clone, Copy, Default)]
pub enum CompressionCodec {
#[default]
/// No compression
None,
/// LZ4 single compression frame with content size present
Lz4,
/// LZ4 raw block compression (no frame). Used by Parquet as `lz4_raw`.
Lz4Raw,
/// Zstandard single compression frame with content size present.
/// Level range is 0–22, where 0 means default compression level (not no compression).
/// Use [`CompressionCodec::zstd_default`] to construct with the default level.
Zstd(u8),
/// Gzip compression. Level range is 0–9, where 0 means no compression.
/// Use [`CompressionCodec::gzip_default`] to construct with the default level.
Gzip(u8),
/// Brotli compression. Level range is 0–11.
/// Use [`CompressionCodec::brotli_default`] to construct with the default level.
Brotli(u8),
/// LZO compression
Lzo,
/// Snappy compression
Snappy,
}
impl CompressionCodec {
/// Returns a Zstd codec with the default compression level.
pub const fn zstd_default() -> Self {
CompressionCodec::Zstd(ZSTD_DEFAULT_LEVEL)
}
/// Returns a Gzip codec with the default compression level.
pub const fn gzip_default() -> Self {
CompressionCodec::Gzip(GZIP_DEFAULT_LEVEL)
}
/// Returns a Brotli codec with the default compression level.
pub const fn brotli_default() -> Self {
CompressionCodec::Brotli(BROTLI_DEFAULT_LEVEL)
}
/// Returns the codec name as used in serialization and error messages.
pub fn name(&self) -> &'static str {
match self {
CompressionCodec::None => "none",
CompressionCodec::Lz4 => "lz4",
CompressionCodec::Lz4Raw => "lz4_raw",
CompressionCodec::Zstd(_) => "zstd",
CompressionCodec::Gzip(_) => "gzip",
CompressionCodec::Brotli(_) => "brotli",
CompressionCodec::Lzo => "lzo",
CompressionCodec::Snappy => "snappy",
}
}
}
// Note: serialize/deserialize do not round-trip the compression level. Iceberg configuration
// only the codec name (e.g. "zstd"), not the level, so deserialization always produces the
// default level. A `Zstd(5)` written to metadata will be read back as `Zstd(3)`. Some
// compression configuration (e.g. Avro metadata) has a separate level field alongside the codec name.
impl Serialize for CompressionCodec {
fn serialize<S: Serializer>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> {
serializer.serialize_str(self.name())
}
}
impl<'de> Deserialize<'de> for CompressionCodec {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> std::result::Result<Self, D::Error> {
let s = String::deserialize(deserializer)?;
match s.to_lowercase().as_str() {
"none" | "uncompressed" => Ok(CompressionCodec::None),
"lz4" => Ok(CompressionCodec::Lz4),
"lz4_raw" => Ok(CompressionCodec::Lz4Raw),
"zstd" => Ok(CompressionCodec::zstd_default()),
"gzip" => Ok(CompressionCodec::gzip_default()),
"brotli" => Ok(CompressionCodec::brotli_default()),
"lzo" => Ok(CompressionCodec::Lzo),
"snappy" => Ok(CompressionCodec::Snappy),
other => Err(serde::de::Error::unknown_variant(other, &[
"none",
"uncompressed",
"lz4",
"lz4_raw",
"zstd",
"gzip",
"brotli",
"lzo",
"snappy",
])),
}
}
}
impl fmt::Display for CompressionCodec {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
CompressionCodec::None => write!(f, "None"),
CompressionCodec::Lz4 => write!(f, "Lz4"),
CompressionCodec::Lz4Raw => write!(f, "Lz4Raw"),
CompressionCodec::Zstd(level) => write!(f, "Zstd(level={level})"),
CompressionCodec::Gzip(level) => write!(f, "Gzip(level={level})"),
CompressionCodec::Brotli(level) => write!(f, "Brotli(level={level})"),
CompressionCodec::Lzo => write!(f, "Lzo"),
CompressionCodec::Snappy => write!(f, "Snappy"),
}
}
}
impl CompressionCodec {
pub(crate) fn decompress(&self, bytes: Vec<u8>) -> Result<Vec<u8>> {
match self {
CompressionCodec::None => Ok(bytes),
CompressionCodec::Lz4 => Err(Error::new(
ErrorKind::FeatureUnsupported,
"LZ4 decompression is not supported currently",
)),
CompressionCodec::Zstd(_) => Ok(zstd::stream::decode_all(&bytes[..])?),
CompressionCodec::Gzip(_) => {
let mut decoder = GzDecoder::new(&bytes[..]);
let mut decompressed = Vec::new();
decoder.read_to_end(&mut decompressed)?;
Ok(decompressed)
}
CompressionCodec::Lz4Raw => Err(Error::new(
ErrorKind::FeatureUnsupported,
"LZ4_RAW decompression is not supported currently",
)),
CompressionCodec::Brotli(_) => Err(Error::new(
ErrorKind::FeatureUnsupported,
"Brotli decompression is not supported currently",
)),
CompressionCodec::Lzo => Err(Error::new(
ErrorKind::FeatureUnsupported,
"LZO decompression is not supported currently",
)),
CompressionCodec::Snappy => Err(Error::new(
ErrorKind::FeatureUnsupported,
"Snappy decompression is not supported currently",
)),
}
}
pub(crate) fn compress(&self, bytes: Vec<u8>) -> Result<Vec<u8>> {
match self {
CompressionCodec::None => Ok(bytes),
CompressionCodec::Lz4 => Err(Error::new(
ErrorKind::FeatureUnsupported,
"LZ4 compression is not supported currently",
)),
CompressionCodec::Zstd(level) => {
let writer = Vec::<u8>::new();
let mut encoder = zstd::stream::Encoder::new(writer, *level as i32)?;
encoder.include_checksum(true)?;
encoder.set_pledged_src_size(Some(bytes.len().try_into()?))?;
std::io::copy(&mut &bytes[..], &mut encoder)?;
Ok(encoder.finish()?)
}
CompressionCodec::Gzip(level) => {
let compression = Compression::new((*level).min(GZIP_MAX_LEVEL) as u32);
let mut encoder = GzEncoder::new(Vec::new(), compression);
encoder.write_all(&bytes)?;
Ok(encoder.finish()?)
}
CompressionCodec::Lz4Raw => Err(Error::new(
ErrorKind::FeatureUnsupported,
"LZ4_RAW compression is not supported currently",
)),
CompressionCodec::Brotli(_) => Err(Error::new(
ErrorKind::FeatureUnsupported,
"Brotli compression is not supported currently",
)),
CompressionCodec::Lzo => Err(Error::new(
ErrorKind::FeatureUnsupported,
"LZO compression is not supported currently",
)),
CompressionCodec::Snappy => Err(Error::new(
ErrorKind::FeatureUnsupported,
"Snappy compression is not supported currently",
)),
}
}
pub(crate) fn is_none(&self) -> bool {
matches!(self, CompressionCodec::None)
}
/// Returns the file extension suffix for this compression codec.
/// Returns empty string for None, ".gz" for Gzip.
///
/// # Errors
///
/// Returns an error for Lz4 and Zstd as they are not fully supported.
pub fn suffix(&self) -> Result<&'static str> {
match self {
CompressionCodec::None => Ok(""),
CompressionCodec::Gzip(_) => Ok(".gz"),
codec @ (CompressionCodec::Lz4
| CompressionCodec::Lz4Raw
| CompressionCodec::Zstd(_)
| CompressionCodec::Brotli(_)
| CompressionCodec::Lzo
| CompressionCodec::Snappy) => Err(Error::new(
ErrorKind::FeatureUnsupported,
format!("suffix not defined for {codec:?}"),
)),
}
}
}
#[cfg(test)]
mod tests {
use super::CompressionCodec;
#[tokio::test]
async fn test_compression_codec_none() {
let bytes_vec = [0_u8; 100].to_vec();
let codec = CompressionCodec::None;
let compressed = codec.compress(bytes_vec.clone()).unwrap();
assert_eq!(bytes_vec, compressed);
let decompressed = codec.decompress(compressed).unwrap();
assert_eq!(bytes_vec, decompressed);
}
#[tokio::test]
async fn test_compression_codec_compress() {
let bytes_vec = [0_u8; 100].to_vec();
let compression_codecs = [
CompressionCodec::zstd_default(),
CompressionCodec::gzip_default(),
];
for codec in compression_codecs {
let compressed = codec.compress(bytes_vec.clone()).unwrap();
assert!(compressed.len() < bytes_vec.len());
let decompressed = codec.decompress(compressed).unwrap();
assert_eq!(decompressed, bytes_vec);
}
}
#[tokio::test]
async fn test_compression_codec_unsupported() {
let unsupported_codecs = [
(CompressionCodec::Lz4, "LZ4"),
(CompressionCodec::Snappy, "Snappy"),
];
let bytes_vec = [0_u8; 100].to_vec();
for (codec, name) in unsupported_codecs {
assert_eq!(
codec.compress(bytes_vec.clone()).unwrap_err().to_string(),
format!("FeatureUnsupported => {name} compression is not supported currently"),
);
assert_eq!(
codec.decompress(bytes_vec.clone()).unwrap_err().to_string(),
format!("FeatureUnsupported => {name} decompression is not supported currently"),
);
}
}
#[test]
fn test_suffix() {
assert_eq!(CompressionCodec::None.suffix().unwrap(), "");
assert_eq!(CompressionCodec::gzip_default().suffix().unwrap(), ".gz");
assert!(CompressionCodec::Lz4.suffix().is_err());
assert!(CompressionCodec::zstd_default().suffix().is_err());
assert!(CompressionCodec::Snappy.suffix().is_err());
let lz4_err = CompressionCodec::Lz4.suffix().unwrap_err();
assert!(lz4_err.to_string().contains("suffix not defined for Lz4"));
let zstd_err = CompressionCodec::zstd_default().suffix().unwrap_err();
assert!(zstd_err.to_string().contains("suffix not defined for Zstd"));
}
#[test]
fn test_display() {
assert_eq!(CompressionCodec::None.to_string(), "None");
assert_eq!(CompressionCodec::Lz4.to_string(), "Lz4");
assert_eq!(
CompressionCodec::zstd_default().to_string(),
"Zstd(level=3)"
);
assert_eq!(CompressionCodec::Zstd(5).to_string(), "Zstd(level=5)");
assert_eq!(
CompressionCodec::gzip_default().to_string(),
"Gzip(level=6)"
);
assert_eq!(CompressionCodec::Gzip(9).to_string(), "Gzip(level=9)");
assert_eq!(CompressionCodec::Snappy.to_string(), "Snappy");
}
}