diff --git a/Cargo.lock b/Cargo.lock index 3802cc03ff..6564f3ce5f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3791,6 +3791,7 @@ dependencies = [ "iceberg-property-macro", "iceberg_test_utils", "itertools 0.13.0", + "lz4_flex", "minijinja", "mockall", "moka", diff --git a/Cargo.toml b/Cargo.toml index 32667b20c7..15eb8de2fd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -115,6 +115,7 @@ itertools = "0.13" libtest-mimic = "0.8.1" linkedbytes = "0.1.8" log = "0.4.28" +lz4_flex = "0.13" metainfo = "0.7.14" mimalloc = "0.1.46" minijinja = "2.12.0" diff --git a/crates/iceberg/Cargo.toml b/crates/iceberg/Cargo.toml index 56e11b1804..69d8385cfd 100644 --- a/crates/iceberg/Cargo.toml +++ b/crates/iceberg/Cargo.toml @@ -62,6 +62,7 @@ form_urlencoded = { workspace = true } futures = { workspace = true } iceberg-property-macro = { workspace = true } itertools = { workspace = true } +lz4_flex = { workspace = true } moka = { version = "0.12.10", features = ["future"] } murmur3 = { workspace = true } once_cell = { workspace = true } diff --git a/crates/iceberg/src/compression.rs b/crates/iceberg/src/compression.rs index dbd2c97881..1cc8e1f346 100644 --- a/crates/iceberg/src/compression.rs +++ b/crates/iceberg/src/compression.rs @@ -23,6 +23,7 @@ use std::io::{Read, Write}; use flate2::Compression; use flate2::read::GzDecoder; use flate2::write::GzEncoder; +use lz4_flex::frame::{FrameDecoder, FrameEncoder, FrameInfo}; use serde::{Deserialize, Deserializer, Serialize, Serializer}; use crate::{Error, ErrorKind, Result}; @@ -149,10 +150,12 @@ impl CompressionCodec { pub(crate) fn decompress(&self, bytes: Vec) -> Result> { match self { CompressionCodec::None => Ok(bytes), - CompressionCodec::Lz4 => Err(Error::new( - ErrorKind::FeatureUnsupported, - "LZ4 decompression is not supported currently", - )), + CompressionCodec::Lz4 => { + let mut decoder = FrameDecoder::new(&bytes[..]); + let mut decompressed = Vec::new(); + decoder.read_to_end(&mut decompressed)?; + Ok(decompressed) + } CompressionCodec::Zstd(_) => Ok(zstd::stream::decode_all(&bytes[..])?), CompressionCodec::Gzip(_) => { let mut decoder = GzDecoder::new(&bytes[..]); @@ -182,10 +185,16 @@ impl CompressionCodec { pub(crate) fn compress(&self, bytes: Vec) -> Result> { match self { CompressionCodec::None => Ok(bytes), - CompressionCodec::Lz4 => Err(Error::new( - ErrorKind::FeatureUnsupported, - "LZ4 compression is not supported currently", - )), + CompressionCodec::Lz4 => { + // Puffin requires one LZ4 frame with content size present: + // https://iceberg.apache.org/puffin-spec/#footer-payload + let frame_info = FrameInfo::new().content_size(Some(bytes.len() as u64)); + let mut encoder = FrameEncoder::with_frame_info(frame_info, Vec::new()); + encoder.write_all(&bytes)?; + encoder.finish().map_err(|e| { + Error::new(ErrorKind::Unexpected, "Failed to finish LZ4 frame").with_source(e) + }) + } CompressionCodec::Zstd(level) => { let writer = Vec::::new(); let mut encoder = zstd::stream::Encoder::new(writer, *level as i32)?; @@ -228,7 +237,7 @@ impl CompressionCodec { /// /// # Errors /// - /// Returns an error for Lz4 and Zstd as they are not fully supported. + /// Returns an error for codecs without a file extension suffix. pub fn suffix(&self) -> Result<&'static str> { match self { CompressionCodec::None => Ok(""), @@ -263,40 +272,53 @@ mod tests { #[tokio::test] async fn test_compression_codec_compress() { - let bytes_vec = [0_u8; 100].to_vec(); - let compression_codecs = [ + CompressionCodec::Lz4, 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); + let empty: Vec = vec![]; + let compressed = codec.compress(empty.clone()).unwrap(); + assert_eq!(codec.decompress(compressed).unwrap(), empty); + + let zeros = vec![0_u8; 100]; + let compressed = codec.compress(zeros.clone()).unwrap(); + assert!(compressed.len() < zeros.len()); + assert_eq!(codec.decompress(compressed).unwrap(), zeros); + + let mixed: Vec = (0..10_000).map(|i| (i % 251) as u8).collect(); + let compressed = codec.compress(mixed.clone()).unwrap(); + assert_eq!(codec.decompress(compressed).unwrap(), mixed); } } #[tokio::test] - async fn test_compression_codec_unsupported() { - let unsupported_codecs = [ - (CompressionCodec::Lz4, "LZ4"), - (CompressionCodec::Snappy, "Snappy"), - ]; + async fn test_lz4_frame_magic_number() { + let compressed = CompressionCodec::Lz4.compress(vec![0u8; 10_000]).unwrap(); + assert_eq!(&compressed[..4], &[0x04, 0x22, 0x4D, 0x18]); + } + + #[tokio::test] + async fn test_snappy_compression_is_unsupported() { 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!( + CompressionCodec::Snappy + .compress(bytes_vec.clone()) + .unwrap_err() + .to_string(), + "FeatureUnsupported => Snappy compression is not supported currently", + ); - assert_eq!( - codec.decompress(bytes_vec.clone()).unwrap_err().to_string(), - format!("FeatureUnsupported => {name} decompression is not supported currently"), - ); - } + assert_eq!( + CompressionCodec::Snappy + .decompress(bytes_vec) + .unwrap_err() + .to_string(), + "FeatureUnsupported => Snappy decompression is not supported currently", + ); } #[test] diff --git a/crates/iceberg/src/puffin/metadata.rs b/crates/iceberg/src/puffin/metadata.rs index 1ee954b873..5959c62ae0 100644 --- a/crates/iceberg/src/puffin/metadata.rs +++ b/crates/iceberg/src/puffin/metadata.rs @@ -570,26 +570,27 @@ mod tests { } #[tokio::test] - async fn test_lz4_compressed_footer_returns_error() { + async fn test_lz4_compressed_footer_is_decoded() { let temp_dir = TempDir::new().unwrap(); + let compressed_payload = CompressionCodec::Lz4 + .compress(empty_footer_payload_bytes()) + .unwrap(); + let mut bytes = vec![]; bytes.extend(FileMetadata::MAGIC.to_vec()); bytes.extend(FileMetadata::MAGIC.to_vec()); - bytes.extend(empty_footer_payload_bytes()); - bytes.extend(empty_footer_payload_bytes_length_bytes()); + bytes.extend(&compressed_payload); + bytes.extend(u32::to_le_bytes(compressed_payload.len() as u32)); bytes.extend(vec![0b00000001, 0, 0, 0]); bytes.extend(FileMetadata::MAGIC.to_vec()); let input_file = input_file_with_bytes(&temp_dir, &bytes).await; assert_eq!( - read_file_metadata(&input_file) - .await - .unwrap_err() - .to_string(), - "FeatureUnsupported => LZ4 decompression is not supported currently", - ) + read_file_metadata(&input_file).await.unwrap(), + empty_footer_payload() + ); } #[tokio::test] diff --git a/crates/iceberg/src/puffin/writer.rs b/crates/iceberg/src/puffin/writer.rs index 0437bd5bf0..bbaa62eb96 100644 --- a/crates/iceberg/src/puffin/writer.rs +++ b/crates/iceberg/src/puffin/writer.rs @@ -303,13 +303,60 @@ mod tests { let blobs = vec![blob_0(), blob_1()]; let blobs_with_compression = blobs_with_compression(blobs.clone(), CompressionCodec::Lz4); - assert_eq!( - write_puffin_file(&temp_dir, blobs_with_compression, file_properties()) - .await - .unwrap_err() - .to_string(), - "FeatureUnsupported => LZ4 compression is not supported currently" - ); + let input_file = write_puffin_file(&temp_dir, blobs_with_compression, file_properties()) + .await + .unwrap() + .to_input_file(); + + assert_eq!(read_all_blobs_from_puffin_file(input_file).await, blobs); + } + + #[tokio::test] + async fn test_compress_footer_lz4_round_trips() { + let temp_dir = TempDir::new().unwrap(); + let file_io = FileIO::new_with_fs(); + let path = temp_dir.path().join("compressed_footer.bin"); + let output_file = file_io.new_output(path.to_str().unwrap()).unwrap(); + + let mut writer = PuffinWriter::new(&output_file, file_properties(), true) + .await + .unwrap(); + writer.add(blob_0(), CompressionCodec::None).await.unwrap(); + writer.close().await.unwrap(); + + let input_file = output_file.to_input_file(); + let bytes = input_file.read().await.unwrap(); + let footer_payload_offset = FileMetadata::MAGIC_LENGTH as usize + + blob_0().data.len() + + FileMetadata::MAGIC_LENGTH as usize; + assert_eq!(&bytes[footer_payload_offset..footer_payload_offset + 4], &[ + 0x04, 0x22, 0x4D, 0x18 + ]); + + let metadata = read_file_metadata(&input_file).await.unwrap(); + assert_eq!(metadata.properties, file_properties()); + assert_eq!(metadata.blobs.len(), 1); + assert_eq!(read_all_blobs_from_puffin_file(input_file).await, vec![ + blob_0() + ]); + } + + #[tokio::test] + async fn test_compress_empty_footer_lz4_succeeds() { + let temp_dir = TempDir::new().unwrap(); + let file_io = FileIO::new_with_fs(); + let path = temp_dir.path().join("compressed_empty_footer.bin"); + let output_file = file_io.new_output(path.to_str().unwrap()).unwrap(); + + let writer = PuffinWriter::new(&output_file, HashMap::new(), true) + .await + .unwrap(); + writer.close().await.unwrap(); + + let input_file = output_file.to_input_file(); + let metadata = read_file_metadata(&input_file).await.unwrap(); + assert!(metadata.blobs.is_empty()); + assert!(metadata.properties.is_empty()); } async fn get_file_as_byte_vec(input_file: InputFile) -> Vec {