From c0ca9c30f9f80532c00c4d14db805540216a83c7 Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Wed, 12 Nov 2025 23:00:56 +0000 Subject: [PATCH] feat!: Support compression codecs for JSON metadata and Avro Previously these properties where not honored on tabel properties. - Adds table properties for these values. - Plumbs them through for writers. --- Cargo.lock | 1 + Cargo.toml | 1 + crates/iceberg/Cargo.toml | 1 + crates/iceberg/public-api.txt | 14 +- crates/iceberg/src/compression.rs | 60 ++++- crates/iceberg/src/io/object_cache.rs | 15 +- crates/iceberg/src/spec/avro_util.rs | 205 ++++++++++++++ crates/iceberg/src/spec/manifest/mod.rs | 13 + crates/iceberg/src/spec/manifest/reader.rs | 4 + crates/iceberg/src/spec/manifest/writer.rs | 252 +++++++++++++++++- .../src/spec/manifest_list/manifest_file.rs | 13 +- crates/iceberg/src/spec/manifest_list/mod.rs | 13 +- .../iceberg/src/spec/manifest_list/writer.rs | 211 ++++++++++++--- crates/iceberg/src/spec/mod.rs | 1 + crates/iceberg/src/spec/table_properties.rs | 129 ++++++++- crates/iceberg/src/test_utils/scan.rs | 47 +++- crates/iceberg/src/transaction/append.rs | 9 +- crates/iceberg/src/transaction/snapshot.rs | 35 ++- 18 files changed, 933 insertions(+), 91 deletions(-) create mode 100644 crates/iceberg/src/spec/avro_util.rs diff --git a/Cargo.lock b/Cargo.lock index e2900acb2e..5adf555c92 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2729,6 +2729,7 @@ dependencies = [ "iceberg_test_utils", "itertools 0.13.0", "minijinja", + "miniz_oxide", "mockall", "moka", "murmur3", diff --git a/Cargo.toml b/Cargo.toml index 661c56b094..4d85923a32 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -100,6 +100,7 @@ itertools = "0.13" linkedbytes = "0.1.8" metainfo = "0.7.14" minijinja = "2.12.0" +miniz_oxide = "0.8" mockall = "0.13.1" mockito = "1" motore-macros = "0.4.3" diff --git a/crates/iceberg/Cargo.toml b/crates/iceberg/Cargo.toml index 8783251fd5..73c2b6f4e8 100644 --- a/crates/iceberg/Cargo.toml +++ b/crates/iceberg/Cargo.toml @@ -61,6 +61,7 @@ form_urlencoded = { workspace = true } futures = { workspace = true } iceberg-property-macro = { workspace = true } itertools = { workspace = true } +miniz_oxide = { workspace = true } moka = { version = "0.12.10", features = ["future"] } murmur3 = { workspace = true } once_cell = { workspace = true } diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index bb4d350e5a..1409513f0a 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -2194,9 +2194,9 @@ impl iceberg::spec::ManifestListWriter pub fn iceberg::spec::ManifestListWriter::add_manifests(&mut self, manifests: impl core::iter::traits::iterator::Iterator) -> iceberg::Result<()> pub async fn iceberg::spec::ManifestListWriter::close(self) -> iceberg::Result pub fn iceberg::spec::ManifestListWriter::next_row_id(&self) -> core::option::Option -pub fn iceberg::spec::ManifestListWriter::v1(writer: alloc::boxed::Box, snapshot_id: i64, parent_snapshot_id: core::option::Option) -> Self -pub fn iceberg::spec::ManifestListWriter::v2(writer: alloc::boxed::Box, snapshot_id: i64, parent_snapshot_id: core::option::Option, sequence_number: i64) -> Self -pub fn iceberg::spec::ManifestListWriter::v3(writer: alloc::boxed::Box, snapshot_id: i64, parent_snapshot_id: core::option::Option, sequence_number: i64, first_row_id: core::option::Option) -> Self +pub fn iceberg::spec::ManifestListWriter::v1(writer: alloc::boxed::Box, snapshot_id: i64, parent_snapshot_id: core::option::Option, compression: iceberg::compression::CompressionCodec) -> iceberg::Result +pub fn iceberg::spec::ManifestListWriter::v2(writer: alloc::boxed::Box, snapshot_id: i64, parent_snapshot_id: core::option::Option, sequence_number: i64, compression: iceberg::compression::CompressionCodec) -> iceberg::Result +pub fn iceberg::spec::ManifestListWriter::v3(writer: alloc::boxed::Box, snapshot_id: i64, parent_snapshot_id: core::option::Option, sequence_number: i64, first_row_id: core::option::Option, compression: iceberg::compression::CompressionCodec) -> iceberg::Result impl core::fmt::Debug for iceberg::spec::ManifestListWriter pub fn iceberg::spec::ManifestListWriter::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result pub struct iceberg::spec::ManifestMetadata @@ -2238,8 +2238,8 @@ pub fn iceberg::spec::ManifestWriterBuilder::build_v2_data(self) -> iceberg::spe pub fn iceberg::spec::ManifestWriterBuilder::build_v2_deletes(self) -> iceberg::spec::ManifestWriter pub fn iceberg::spec::ManifestWriterBuilder::build_v3_data(self) -> iceberg::spec::ManifestWriter pub fn iceberg::spec::ManifestWriterBuilder::build_v3_deletes(self) -> iceberg::spec::ManifestWriter -pub fn iceberg::spec::ManifestWriterBuilder::new(output: iceberg::io::OutputFile, snapshot_id: core::option::Option, schema: iceberg::spec::SchemaRef, partition_spec: iceberg::spec::PartitionSpec) -> Self -pub fn iceberg::spec::ManifestWriterBuilder::new_from_encrypted(encrypted_output: iceberg::encryption::EncryptedOutputFile, snapshot_id: core::option::Option, schema: iceberg::spec::SchemaRef, partition_spec: iceberg::spec::PartitionSpec) -> iceberg::Result +pub fn iceberg::spec::ManifestWriterBuilder::new(output: iceberg::io::OutputFile, snapshot_id: core::option::Option, schema: iceberg::spec::SchemaRef, partition_spec: iceberg::spec::PartitionSpec, compression: iceberg::compression::CompressionCodec) -> iceberg::Result +pub fn iceberg::spec::ManifestWriterBuilder::new_from_encrypted(encrypted_output: iceberg::encryption::EncryptedOutputFile, snapshot_id: core::option::Option, schema: iceberg::spec::SchemaRef, partition_spec: iceberg::spec::PartitionSpec, compression: iceberg::compression::CompressionCodec) -> iceberg::Result pub struct iceberg::spec::Map impl iceberg::spec::Map pub fn iceberg::spec::Map::get(&self, key: &iceberg::spec::Literal) -> core::option::Option<&core::option::Option> @@ -2862,6 +2862,9 @@ impl core::fmt::Debug for iceberg::spec::TableMetadataBuilder pub fn iceberg::spec::TableMetadataBuilder::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result pub struct iceberg::spec::TableProperties<'properties> impl iceberg::spec::TableProperties<'_> +pub const iceberg::spec::TableProperties<'_>::PROPERTY_AVRO_COMPRESSION_CODEC: &'static str +pub const iceberg::spec::TableProperties<'_>::PROPERTY_AVRO_COMPRESSION_CODEC_DEFAULT: &'static str +pub const iceberg::spec::TableProperties<'_>::PROPERTY_AVRO_COMPRESSION_LEVEL: &'static str pub const iceberg::spec::TableProperties<'_>::PROPERTY_COMMIT_MAX_RETRY_WAIT_MS: &'static str pub const iceberg::spec::TableProperties<'_>::PROPERTY_COMMIT_MAX_RETRY_WAIT_MS_DEFAULT: u64 pub const iceberg::spec::TableProperties<'_>::PROPERTY_COMMIT_MIN_RETRY_WAIT_MS: &'static str @@ -2934,6 +2937,7 @@ pub const iceberg::spec::TableProperties<'_>::PROPERTY_WRITE_TARGET_FILE_SIZE_BY pub const iceberg::spec::TableProperties<'_>::RESERVED_PROPERTIES: [&'static str; 9] pub fn iceberg::spec::TableProperties<'_>::data_encryption_key_size(&self) -> iceberg::Result impl<'properties> iceberg::spec::TableProperties<'properties> +pub fn iceberg::spec::TableProperties<'properties>::avro_compression_codec(&self) -> iceberg::Result pub fn iceberg::spec::TableProperties<'properties>::cdc_enabled(&self) -> iceberg::Result pub fn iceberg::spec::TableProperties<'properties>::cdc_max_chunk_size(&self) -> iceberg::Result pub fn iceberg::spec::TableProperties<'properties>::cdc_min_chunk_size(&self) -> iceberg::Result diff --git a/crates/iceberg/src/compression.rs b/crates/iceberg/src/compression.rs index dbd2c97881..28fa0b570a 100644 --- a/crates/iceberg/src/compression.rs +++ b/crates/iceberg/src/compression.rs @@ -91,6 +91,22 @@ impl CompressionCodec { CompressionCodec::Snappy => "snappy", } } + + /// Parses a codec name (case-insensitive), returning `None` for unrecognized names. + /// Codecs that carry a level get their default level. + pub(crate) fn from_name(name: &str) -> Option { + Some(match name.to_lowercase().as_str() { + "none" | "uncompressed" => CompressionCodec::None, + "lz4" => CompressionCodec::Lz4, + "lz4_raw" => CompressionCodec::Lz4Raw, + "zstd" => CompressionCodec::zstd_default(), + "gzip" => CompressionCodec::gzip_default(), + "brotli" => CompressionCodec::brotli_default(), + "lzo" => CompressionCodec::Lzo, + "snappy" => CompressionCodec::Snappy, + _ => return None, + }) + } } // Note: serialize/deserialize do not round-trip the compression level. Iceberg configuration @@ -106,16 +122,8 @@ impl Serialize for CompressionCodec { impl<'de> Deserialize<'de> for CompressionCodec { fn deserialize>(deserializer: D) -> std::result::Result { 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, &[ + Self::from_name(&s).ok_or_else(|| { + serde::de::Error::unknown_variant(&s, &[ "none", "uncompressed", "lz4", @@ -125,8 +133,8 @@ impl<'de> Deserialize<'de> for CompressionCodec { "brotli", "lzo", "snappy", - ])), - } + ]) + }) } } @@ -219,6 +227,17 @@ impl CompressionCodec { } } + /// Replace the compression level, if this variant carries one. + /// Variants without a level are returned unchanged. + pub(crate) fn with_level(self, level: u8) -> Self { + match self { + CompressionCodec::Zstd(_) => CompressionCodec::Zstd(level), + CompressionCodec::Gzip(_) => CompressionCodec::Gzip(level), + CompressionCodec::Brotli(_) => CompressionCodec::Brotli(level), + other => other, + } + } + pub(crate) fn is_none(&self) -> bool { matches!(self, CompressionCodec::None) } @@ -315,6 +334,23 @@ mod tests { assert!(zstd_err.to_string().contains("suffix not defined for Zstd")); } + #[test] + fn test_with_level() { + for (codec, expected) in [ + (CompressionCodec::zstd_default(), CompressionCodec::Zstd(5)), + (CompressionCodec::gzip_default(), CompressionCodec::Gzip(5)), + ( + CompressionCodec::brotli_default(), + CompressionCodec::Brotli(5), + ), + (CompressionCodec::None, CompressionCodec::None), + (CompressionCodec::Lz4, CompressionCodec::Lz4), + (CompressionCodec::Snappy, CompressionCodec::Snappy), + ] { + assert_eq!(codec.with_level(5), expected); + } + } + #[test] fn test_display() { assert_eq!(CompressionCodec::None.to_string(), "None"); diff --git a/crates/iceberg/src/io/object_cache.rs b/crates/iceberg/src/io/object_cache.rs index dfc6c124c9..6e2334b253 100644 --- a/crates/iceberg/src/io/object_cache.rs +++ b/crates/iceberg/src/io/object_cache.rs @@ -225,6 +225,7 @@ mod tests { use super::*; use crate::TableIdent; + use crate::compression::CompressionCodec; use crate::io::{FileIO, OutputFile}; use crate::spec::{ DataContentType, DataFileBuilder, DataFileFormat, FormatVersion, Literal, @@ -307,7 +308,9 @@ mod tests { Some(current_snapshot.snapshot_id()), current_schema.clone(), current_partition_spec.as_ref().clone(), + CompressionCodec::None, ) + .unwrap() .build_v2_data(); writer .add_entry( @@ -344,7 +347,9 @@ mod tests { current_snapshot.snapshot_id(), current_snapshot.parent_snapshot_id(), current_snapshot.sequence_number(), - ); + CompressionCodec::None, + ) + .unwrap(); manifest_list_write .add_manifests(vec![data_file_manifest].into_iter()) .unwrap(); @@ -362,7 +367,9 @@ mod tests { Some(current_snapshot.snapshot_id()), current_schema.clone(), current_partition_spec.as_ref().clone(), + CompressionCodec::None, ) + .unwrap() .build_v1(); writer .add_entry( @@ -398,7 +405,9 @@ mod tests { manifest_list_writer, current_snapshot.snapshot_id(), current_snapshot.parent_snapshot_id(), - ); + CompressionCodec::None, + ) + .unwrap(); manifest_list_write .add_manifests(vec![data_file_manifest].into_iter()) .unwrap(); @@ -618,7 +627,9 @@ mod tests { Some(1), schema, partition_spec, + CompressionCodec::None, ) + .unwrap() .build_v3_data(); writer .add_entry( diff --git a/crates/iceberg/src/spec/avro_util.rs b/crates/iceberg/src/spec/avro_util.rs new file mode 100644 index 0000000000..9bfb5784e8 --- /dev/null +++ b/crates/iceberg/src/spec/avro_util.rs @@ -0,0 +1,205 @@ +// 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. + +//! Utilities for working with Apache Avro in Iceberg. + +use apache_avro::{Codec, DeflateSettings, ZstandardSettings}; +use miniz_oxide::deflate::CompressionLevel; + +use crate::compression::CompressionCodec; +use crate::error::{Result, invalid_data}; + +// Avro manifest defaults, chosen to match Java Iceberg (GzipCodec/ZstandardCodec). Distinct from +// the data-file defaults in compression.rs (gzip=6, zstd=3). +const DEFAULT_GZIP_LEVEL: u8 = 9; +const DEFAULT_ZSTD_LEVEL: u8 = 1; +/// Max gzip level, matching the range accepted by Java's `Deflater`. +const MAX_GZIP_LEVEL: u8 = 9; +/// Max supported level for ZSTD; higher levels are clamped, as zstd itself does. +const MAX_ZSTD_LEVEL: u8 = 22; + +/// Default codec for Avro files (manifests, manifest lists). +pub(crate) const DEFAULT_AVRO_CODEC: CompressionCodec = CompressionCodec::Gzip(DEFAULT_GZIP_LEVEL); + +/// Parse a codec name (case-insensitive) into a [`CompressionCodec`], with the Avro-specific +/// default level for codecs that carry one. +/// +/// This only resolves the name; use [`to_avro_codec`] to check that the codec (and its level) is +/// supported for Avro. +pub(crate) fn parse_avro_codec(name: &str) -> Result { + Ok( + match CompressionCodec::from_name(name) + .ok_or_else(|| invalid_data!("Unsupported compression codec: {name}"))? + { + CompressionCodec::Gzip(_) => CompressionCodec::Gzip(DEFAULT_GZIP_LEVEL), + CompressionCodec::Zstd(_) => CompressionCodec::Zstd(DEFAULT_ZSTD_LEVEL), + other => other, + }, + ) +} + +/// Convert a [`CompressionCodec`] to an [`apache_avro::Codec`] for use in Avro writers. +/// +/// Returns an error for codecs that are not supported in Avro (e.g. LZ4) and for gzip levels +/// above 9. Matching Java, zstd levels above 22 are clamped. +/// +/// apache-avro only exposes the named deflate levels 0, 1, 6 and 9, so other gzip levels are +/// rounded to the nearest of these (1–3 → 1, 4–7 → 6, 8 → 9). +pub(crate) fn to_avro_codec(codec: CompressionCodec) -> Result { + match codec { + CompressionCodec::None => Ok(Codec::Null), + CompressionCodec::Snappy => Ok(Codec::Snappy), + CompressionCodec::Gzip(level) => { + let compression_level = match level { + 0 => CompressionLevel::NoCompression, + 1..=3 => CompressionLevel::BestSpeed, + 4..=7 => CompressionLevel::DefaultLevel, + 8..=MAX_GZIP_LEVEL => CompressionLevel::BestCompression, + _ => { + return Err(invalid_data!( + "Invalid gzip compression level {level} for Avro files, expected 0-{MAX_GZIP_LEVEL}" + )); + } + }; + Ok(Codec::Deflate(DeflateSettings::new(compression_level))) + } + CompressionCodec::Zstd(level) => Ok(Codec::Zstandard(ZstandardSettings::new( + level.min(MAX_ZSTD_LEVEL), + ))), + CompressionCodec::Lz4 + | CompressionCodec::Lz4Raw + | CompressionCodec::Brotli(_) + | CompressionCodec::Lzo => Err(invalid_data!( + "Unsupported compression codec for Avro files: {}", + codec.name() + )), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ErrorKind; + + fn deflate(level: CompressionLevel) -> Codec { + Codec::Deflate(DeflateSettings::new(level)) + } + + #[test] + fn test_parse_avro_codec() { + for (name, expected) in [ + // Mixed case verifies case-insensitive matching. + ("GZip", CompressionCodec::Gzip(DEFAULT_GZIP_LEVEL)), + ("zstd", CompressionCodec::Zstd(DEFAULT_ZSTD_LEVEL)), + ("snappy", CompressionCodec::Snappy), + ("uncompressed", CompressionCodec::None), + ("none", CompressionCodec::None), + ] { + assert_eq!(parse_avro_codec(name).unwrap(), expected, "codec {name}"); + } + } + + #[test] + fn test_parse_avro_codec_unknown() { + let err = parse_avro_codec("unknown").unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + assert!( + err.to_string() + .contains("Unsupported compression codec: unknown"), + "unexpected error: {err}" + ); + } + + #[test] + fn test_to_avro_codec() { + for (codec, expected) in [ + (CompressionCodec::None, Codec::Null), + (CompressionCodec::Snappy, Codec::Snappy), + ( + DEFAULT_AVRO_CODEC, + deflate(CompressionLevel::BestCompression), + ), + ( + CompressionCodec::Zstd(3), + Codec::Zstandard(ZstandardSettings::new(3)), + ), + ( + CompressionCodec::Zstd(MAX_ZSTD_LEVEL + 1), + Codec::Zstandard(ZstandardSettings::new(MAX_ZSTD_LEVEL)), + ), + ] { + assert_eq!(to_avro_codec(codec).unwrap(), expected, "codec {codec}"); + } + } + + #[test] + fn test_to_avro_codec_gzip_levels_round_to_nearest() { + let expected = [ + CompressionLevel::NoCompression, + CompressionLevel::BestSpeed, + CompressionLevel::BestSpeed, + CompressionLevel::BestSpeed, + CompressionLevel::DefaultLevel, + CompressionLevel::DefaultLevel, + CompressionLevel::DefaultLevel, + CompressionLevel::DefaultLevel, + CompressionLevel::BestCompression, + CompressionLevel::BestCompression, + ]; + for (level, expected) in (0..=MAX_GZIP_LEVEL).zip(expected) { + assert_eq!( + to_avro_codec(CompressionCodec::Gzip(level)).unwrap(), + deflate(expected), + "gzip level {level}" + ); + } + } + + #[test] + fn test_to_avro_codec_rejected() { + for (codec, message) in [ + ( + CompressionCodec::Gzip(MAX_GZIP_LEVEL + 1), + "Invalid gzip compression level 10 for Avro files, expected 0-9", + ), + ( + CompressionCodec::Gzip(200), + "Invalid gzip compression level 200 for Avro files, expected 0-9", + ), + ( + CompressionCodec::Lz4, + "Unsupported compression codec for Avro files: lz4", + ), + ( + CompressionCodec::Lz4Raw, + "Unsupported compression codec for Avro files: lz4_raw", + ), + ( + CompressionCodec::brotli_default(), + "Unsupported compression codec for Avro files: brotli", + ), + ( + CompressionCodec::Lzo, + "Unsupported compression codec for Avro files: lzo", + ), + ] { + let err = to_avro_codec(codec).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid, "codec {codec}"); + assert!(err.to_string().contains(message), "unexpected error: {err}"); + } + } +} diff --git a/crates/iceberg/src/spec/manifest/mod.rs b/crates/iceberg/src/spec/manifest/mod.rs index 3bd4dc8ad8..14ec6b4e57 100644 --- a/crates/iceberg/src/spec/manifest/mod.rs +++ b/crates/iceberg/src/spec/manifest/mod.rs @@ -163,6 +163,7 @@ mod tests { use tempfile::TempDir; use super::*; + use crate::compression::CompressionCodec; use crate::io::FileIO; use crate::spec::{Literal, NestedField, PrimitiveType, Struct, Transform, Type}; @@ -269,7 +270,9 @@ mod tests { Some(1), metadata.schema.clone(), metadata.partition_spec.clone(), + CompressionCodec::None, ) + .unwrap() .build_v2_data(); for entry in &entries { writer.add_entry(entry.clone()).unwrap(); @@ -574,7 +577,9 @@ mod tests { Some(2), metadata.schema.clone(), metadata.partition_spec.clone(), + CompressionCodec::None, ) + .unwrap() .build_v2_data(); for entry in &entries { writer.add_entry(entry.clone()).unwrap(); @@ -670,7 +675,9 @@ mod tests { Some(3), metadata.schema.clone(), metadata.partition_spec.clone(), + CompressionCodec::None, ) + .unwrap() .build_v1(); for entry in &entries { writer.add_entry(entry.clone()).unwrap(); @@ -778,7 +785,9 @@ mod tests { Some(2), metadata.schema.clone(), metadata.partition_spec.clone(), + CompressionCodec::None, ) + .unwrap() .build_v1(); for entry in &entries { writer.add_entry(entry.clone()).unwrap(); @@ -885,7 +894,9 @@ mod tests { Some(2), metadata.schema.clone(), metadata.partition_spec.clone(), + CompressionCodec::None, ) + .unwrap() .build_v2_data(); for entry in &entries { writer.add_entry(entry.clone()).unwrap(); @@ -1163,7 +1174,9 @@ mod tests { Some(1), metadata.schema.clone(), metadata.partition_spec.clone(), + CompressionCodec::None, ) + .unwrap() .build_v2_data(); for entry in &entries { writer.add_entry(entry.clone()).unwrap(); diff --git a/crates/iceberg/src/spec/manifest/reader.rs b/crates/iceberg/src/spec/manifest/reader.rs index e7044a56bc..0d68804553 100644 --- a/crates/iceberg/src/spec/manifest/reader.rs +++ b/crates/iceberg/src/spec/manifest/reader.rs @@ -132,6 +132,7 @@ mod tests { use super::*; use crate::ErrorKind; + use crate::compression::CompressionCodec; use crate::encryption::{EncryptedOutputFile, StandardKeyMetadata}; use crate::io::FileIO; use crate::spec::{ @@ -154,7 +155,9 @@ mod tests { Some(1), schema.clone(), partition_spec, + CompressionCodec::None, ) + .unwrap() .build_v2_data(); writer.add_entry(test_entry()).unwrap(); // Writing the manifest yields the manifest list entry describing it. @@ -195,6 +198,7 @@ mod tests { Some(1), schema.clone(), partition_spec, + CompressionCodec::None, ) .unwrap() .build_v3_data(); diff --git a/crates/iceberg/src/spec/manifest/writer.rs b/crates/iceberg/src/spec/manifest/writer.rs index 15608a63f8..0946f4f95f 100644 --- a/crates/iceberg/src/spec/manifest/writer.rs +++ b/crates/iceberg/src/spec/manifest/writer.rs @@ -17,7 +17,7 @@ use std::cmp::min; -use apache_avro::{Writer as AvroWriter, to_value}; +use apache_avro::{Codec, Writer as AvroWriter, to_value}; use bytes::Bytes; use itertools::Itertools; use serde_json::to_vec; @@ -26,6 +26,7 @@ use super::{ Datum, FormatVersion, ManifestContentType, PartitionSpec, PrimitiveType, UNASSIGNED_SEQUENCE_NUMBER, }; +use crate::compression::CompressionCodec; use crate::encryption::EncryptedOutputFile; use crate::error::{Result, invalid_data}; use crate::io::{FileMetadata, FileWrite, OutputFile}; @@ -33,7 +34,7 @@ use crate::spec::manifest::_serde::{ManifestEntryV1, ManifestEntryV2}; use crate::spec::manifest::{manifest_schema_v1, manifest_schema_v2}; use crate::spec::{ DataContentType, DataFile, FieldSummary, ManifestEntry, ManifestFile, ManifestMetadata, - ManifestStatus, PrimitiveLiteral, SchemaRef, StructType, Type, + ManifestStatus, PrimitiveLiteral, SchemaRef, StructType, Type, avro_util, }; /// Placeholder for snapshot ID. The field with this value must be replaced @@ -74,42 +75,68 @@ pub struct ManifestWriterBuilder { snapshot_id: Option, schema: SchemaRef, partition_spec: PartitionSpec, + codec: Codec, } impl ManifestWriterBuilder { /// Create a new builder for unencrypted manifests. + /// + /// Returns an error if `compression` is not supported for Avro files. pub fn new( output: OutputFile, snapshot_id: Option, schema: SchemaRef, partition_spec: PartitionSpec, - ) -> Self { + compression: CompressionCodec, + ) -> Result { let location = output.location().to_owned(); - Self { - output: ManifestOutput::Plain(output), + Self::from_output( + ManifestOutput::Plain(output), location, snapshot_id, schema, partition_spec, - } + compression, + ) } /// Create a new builder from an [`EncryptedOutputFile`]. /// /// Use this when writing manifests with transparent encryption. + /// Returns an error if `compression` is not supported for Avro files. pub fn new_from_encrypted( encrypted_output: EncryptedOutputFile, snapshot_id: Option, schema: SchemaRef, partition_spec: PartitionSpec, + compression: CompressionCodec, ) -> Result { let location = encrypted_output.location().to_owned(); + Self::from_output( + ManifestOutput::Encrypted(encrypted_output), + location, + snapshot_id, + schema, + partition_spec, + compression, + ) + } + + fn from_output( + output: ManifestOutput, + location: String, + snapshot_id: Option, + schema: SchemaRef, + partition_spec: PartitionSpec, + compression: CompressionCodec, + ) -> Result { Ok(Self { - output: ManifestOutput::Encrypted(encrypted_output), + output, location, snapshot_id, schema, partition_spec, + codec: avro_util::to_avro_codec(compression)?, }) } @@ -122,7 +149,14 @@ impl ManifestWriterBuilder { .format_version(FormatVersion::V1) .content(ManifestContentType::Data) .build(); - ManifestWriter::new(self.output, self.location, self.snapshot_id, metadata, None) + ManifestWriter::new( + self.output, + self.location, + self.snapshot_id, + metadata, + None, + self.codec, + ) } /// Build a [`ManifestWriter`] for format version 2, data content. @@ -135,7 +169,14 @@ impl ManifestWriterBuilder { .content(ManifestContentType::Data) .build(); - ManifestWriter::new(self.output, self.location, self.snapshot_id, metadata, None) + ManifestWriter::new( + self.output, + self.location, + self.snapshot_id, + metadata, + None, + self.codec, + ) } /// Build a [`ManifestWriter`] for format version 2, deletes content. @@ -147,7 +188,14 @@ impl ManifestWriterBuilder { .format_version(FormatVersion::V2) .content(ManifestContentType::Deletes) .build(); - ManifestWriter::new(self.output, self.location, self.snapshot_id, metadata, None) + ManifestWriter::new( + self.output, + self.location, + self.snapshot_id, + metadata, + None, + self.codec, + ) } /// Build a [`ManifestWriter`] for format version 2, data content. @@ -167,6 +215,7 @@ impl ManifestWriterBuilder { // First row id is assigned by the [`ManifestListWriter`] when the manifest // is added to the list. None, + self.codec, ) } @@ -179,7 +228,14 @@ impl ManifestWriterBuilder { .format_version(FormatVersion::V3) .content(ManifestContentType::Deletes) .build(); - ManifestWriter::new(self.output, self.location, self.snapshot_id, metadata, None) + ManifestWriter::new( + self.output, + self.location, + self.snapshot_id, + metadata, + None, + self.codec, + ) } } @@ -203,6 +259,8 @@ pub struct ManifestWriter { manifest_entries: Vec, metadata: ManifestMetadata, + + codec: Codec, } impl ManifestWriter { @@ -213,6 +271,7 @@ impl ManifestWriter { snapshot_id: Option, metadata: ManifestMetadata, first_row_id: Option, + codec: Codec, ) -> Self { Self { output, @@ -228,6 +287,7 @@ impl ManifestWriter { min_seq_num: None, manifest_entries: Vec::new(), metadata, + codec, } } @@ -432,7 +492,8 @@ impl ManifestWriter { // Manifest schema did not change between V2 and V3 FormatVersion::V2 | FormatVersion::V3 => manifest_schema_v2(&partition_type)?, }; - let mut avro_writer = AvroWriter::new(&avro_schema, Vec::new()); + + let mut avro_writer = AvroWriter::with_codec(&avro_schema, Vec::new(), self.codec); avro_writer.add_user_metadata( "schema".to_string(), to_vec(table_schema) @@ -575,13 +636,19 @@ impl PartitionFieldStats { mod tests { use std::collections::HashMap; use std::fs; + use std::path::Path; use std::sync::Arc; use tempfile::TempDir; use super::*; + use crate::ErrorKind; + use crate::encryption::StandardKeyMetadata; use crate::io::FileIO; - use crate::spec::{DataFileFormat, Manifest, NestedField, PrimitiveType, Schema, Struct, Type}; + use crate::spec::{ + DataFileBuilder, DataFileFormat, Manifest, ManifestReader, NestedField, PrimitiveType, + Schema, Struct, Type, + }; #[tokio::test] async fn test_add_delete_existing() { @@ -712,7 +779,9 @@ mod tests { Some(3), metadata.schema.clone(), metadata.partition_spec.clone(), + CompressionCodec::None, ) + .unwrap() .build_v2_data(); writer.add_entry(entries[0].clone()).unwrap(); writer.add_delete_entry(entries[1].clone()).unwrap(); @@ -732,6 +801,116 @@ mod tests { assert_eq!(actual_manifest, Manifest::new(metadata, entries)); } + fn added_entry(snapshot_id: i64, path: String, size: u64, records: u64) -> ManifestEntry { + let data_file = DataFileBuilder::default() + .content(DataContentType::Data) + .file_path(path) + .file_format(DataFileFormat::Parquet) + .partition(Struct::empty()) + .file_size_in_bytes(size) + .record_count(records) + .build() + .unwrap(); + ManifestEntry::builder() + .status(ManifestStatus::Added) + .snapshot_id(snapshot_id) + .sequence_number(1) + .file_sequence_number(1) + .data_file(data_file) + .build() + } + + #[test] + fn test_manifest_writer_builder_rejects_unsupported_codec() { + let io = FileIO::new_with_memory(); + let output = io.new_output("memory:///test/lz4_manifest.avro").unwrap(); + let err = ManifestWriterBuilder::new( + output, + Some(1), + Arc::new(Schema::builder().build().unwrap()), + PartitionSpec::unpartition_spec(), + CompressionCodec::Lz4, + ) + .err() + .unwrap(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + assert!(err.to_string().contains("lz4"), "unexpected error: {err}"); + } + + #[tokio::test] + async fn test_manifest_writer_with_compression() { + let metadata = { + let schema = Schema::builder() + .with_fields(vec![Arc::new(NestedField::required( + 1, + "id", + Type::Primitive(PrimitiveType::Int), + ))]) + .build() + .unwrap(); + + ManifestMetadata { + schema_id: 0, + schema: Arc::new(schema), + partition_spec: PartitionSpec::unpartition_spec(), + format_version: FormatVersion::V2, + content: ManifestContentType::Data, + } + }; + + async fn write_manifest( + io: &FileIO, + path: &Path, + metadata: &ManifestMetadata, + compression: CompressionCodec, + ) { + let output_file = io.new_output(path.to_str().unwrap()).unwrap(); + let mut writer = ManifestWriterBuilder::new( + output_file, + Some(1), + metadata.schema.clone(), + metadata.partition_spec.clone(), + compression, + ) + .unwrap() + .build_v2_data(); + for i in 0..1000 { + let entry = added_entry( + 1, + format!( + "/very/long/path/to/data/directory/with/many/subdirectories/file_{i}.parquet" + ), + 100000 + i, + 1000 + i, + ); + writer.add_entry(entry).unwrap(); + } + writer.write_manifest_file().await.unwrap(); + } + + let tmp_dir = TempDir::new().unwrap(); + let io = FileIO::new_with_fs(); + let uncompressed_path = tmp_dir.path().join("uncompressed_manifest.avro"); + let compressed_path = tmp_dir.path().join("compressed_manifest.avro"); + + write_manifest(&io, &uncompressed_path, &metadata, CompressionCodec::None).await; + write_manifest(&io, &compressed_path, &metadata, CompressionCodec::Gzip(9)).await; + + let uncompressed_size = fs::metadata(&uncompressed_path).unwrap().len(); + let compressed_size = fs::metadata(&compressed_path).unwrap().len(); + + assert!( + compressed_size < uncompressed_size, + "Compressed size ({compressed_size}) should be less than uncompressed size ({uncompressed_size})" + ); + + // Verify the compressed file can be read back correctly + let compressed_bytes = fs::read(&compressed_path).unwrap(); + let manifest = Manifest::parse_avro(&compressed_bytes).unwrap(); + assert_eq!(manifest.metadata.format_version, FormatVersion::V2); + assert_eq!(manifest.entries.len(), 1000); + } + #[tokio::test] async fn test_v3_delete_manifest_delete_file_roundtrip() { let schema = Arc::new( @@ -799,7 +978,9 @@ mod tests { Some(1), schema.clone(), partition_spec.clone(), + CompressionCodec::None, ) + .unwrap() .build_v3_deletes(); writer.add_entry(delete_entry).unwrap(); @@ -819,4 +1000,49 @@ mod tests { ManifestContentType::Deletes, ); } + + #[tokio::test] + async fn test_manifest_writer_encrypted_with_compression() { + let schema = Arc::new( + Schema::builder() + .with_fields(vec![Arc::new(NestedField::required( + 1, + "id", + Type::Primitive(PrimitiveType::Int), + ))]) + .build() + .unwrap(), + ); + let partition_spec = PartitionSpec::unpartition_spec(); + let key_metadata = StandardKeyMetadata::try_new(b"0123456789abcdef") + .unwrap() + .with_aad_prefix(b"test-aad-prefix!"); + + let io = FileIO::new_with_memory(); + let path = "memory:///test/encrypted_manifest.avro"; + let output = io.new_output(path).unwrap(); + let encrypted_output = EncryptedOutputFile::new(output, key_metadata); + + let mut writer = ManifestWriterBuilder::new_from_encrypted( + encrypted_output, + Some(42), + schema.clone(), + partition_spec.clone(), + CompressionCodec::Gzip(9), + ) + .unwrap() + .build_v2_data(); + + for i in 0..10 { + let entry = added_entry(42, format!("/data/file_{i}.parquet"), 1000 + i, 100 + i); + writer.add_entry(entry).unwrap(); + } + let manifest_file = writer.write_manifest_file().await.unwrap(); + assert_eq!(manifest_file.added_files_count, Some(10)); + + // Read back via the recorded key metadata and verify the manifest parses correctly. + let manifest = ManifestReader::new(io).read(&manifest_file).await.unwrap(); + assert_eq!(manifest.entries().len(), 10); + assert_eq!(manifest.metadata().format_version, FormatVersion::V2); + } } diff --git a/crates/iceberg/src/spec/manifest_list/manifest_file.rs b/crates/iceberg/src/spec/manifest_list/manifest_file.rs index 24df6d4dca..de595882b3 100644 --- a/crates/iceberg/src/spec/manifest_list/manifest_file.rs +++ b/crates/iceberg/src/spec/manifest_list/manifest_file.rs @@ -197,6 +197,7 @@ mod test { use super::{ManifestContentType, ManifestFile}; use crate::ErrorKind; + use crate::compression::CompressionCodec; use crate::encryption::{EncryptedOutputFile, StandardKeyMetadata}; use crate::io::FileIO; use crate::spec::{ @@ -239,8 +240,15 @@ mod test { .unwrap(); let output_file = io.new_output(path).unwrap(); - let mut writer = ManifestWriterBuilder::new(output_file, Some(1), schema, partition_spec) - .build_v3_data(); + let mut writer = ManifestWriterBuilder::new( + output_file, + Some(1), + schema, + partition_spec, + CompressionCodec::None, + ) + .unwrap() + .build_v3_data(); writer .add_entry(data_entry(ManifestStatus::Added, 100, None)) @@ -270,6 +278,7 @@ mod test { Some(1), schema, partition_spec, + CompressionCodec::None, ) .expect("Expected a valid writer") .build_v3_data(); diff --git a/crates/iceberg/src/spec/manifest_list/mod.rs b/crates/iceberg/src/spec/manifest_list/mod.rs index 25dd5bb93e..93c61cf24f 100644 --- a/crates/iceberg/src/spec/manifest_list/mod.rs +++ b/crates/iceberg/src/spec/manifest_list/mod.rs @@ -99,6 +99,7 @@ mod test { use super::_const_schema::MANIFEST_LIST_AVRO_SCHEMA_V2; use super::_serde::ManifestFileV2; use super::*; + use crate::compression::CompressionCodec; use crate::io::FileIO; use crate::spec::{Datum, FieldSummary, ManifestContentType, ManifestFile}; @@ -142,7 +143,9 @@ mod test { .unwrap(), 1646658105718557341, Some(1646658105718557341), - ); + CompressionCodec::None, + ) + .unwrap(); writer .add_manifests(manifest_list.entries.clone().into_iter()) @@ -220,7 +223,9 @@ mod test { 1646658105718557341, Some(1646658105718557341), 1, - ); + CompressionCodec::None, + ) + .unwrap(); writer .add_manifests(manifest_list.entries.clone().into_iter()) @@ -339,7 +344,9 @@ mod test { Some(377075049360453639), 1, Some(10), - ); + CompressionCodec::None, + ) + .unwrap(); writer .add_manifests(manifest_list.entries.clone().into_iter()) diff --git a/crates/iceberg/src/spec/manifest_list/writer.rs b/crates/iceberg/src/spec/manifest_list/writer.rs index 5ce2cfb259..7e9bbce53a 100644 --- a/crates/iceberg/src/spec/manifest_list/writer.rs +++ b/crates/iceberg/src/spec/manifest_list/writer.rs @@ -25,8 +25,10 @@ use super::_const_schema::{ }; use super::_serde::{ManifestFileV1, ManifestFileV2, ManifestFileV3}; use super::{FormatVersion, ManifestContentType, ManifestFile, UNASSIGNED_SEQUENCE_NUMBER}; +use crate::compression::CompressionCodec; use crate::error::{Result, invalid_data}; use crate::io::{FileMetadata, FileWrite}; +use crate::spec::avro_util::to_avro_codec; use crate::{Error, ErrorKind}; /// A manifest list writer. @@ -59,7 +61,8 @@ impl ManifestListWriter { writer: Box, snapshot_id: i64, parent_snapshot_id: Option, - ) -> Self { + compression: CompressionCodec, + ) -> Result { let mut metadata = HashMap::from_iter([ ("snapshot-id".to_string(), snapshot_id.to_string()), ("format-version".to_string(), "1".to_string()), @@ -70,7 +73,15 @@ impl ManifestListWriter { parent_snapshot_id.to_string(), ); } - Self::new(FormatVersion::V1, writer, metadata, 0, snapshot_id, None) + Self::new( + FormatVersion::V1, + writer, + metadata, + 0, + snapshot_id, + None, + compression, + ) } /// Construct a v2 [`ManifestListWriter`] that writes to a provided [`FileWrite`]. @@ -79,7 +90,8 @@ impl ManifestListWriter { snapshot_id: i64, parent_snapshot_id: Option, sequence_number: i64, - ) -> Self { + compression: CompressionCodec, + ) -> Result { let mut metadata = HashMap::from_iter([ ("snapshot-id".to_string(), snapshot_id.to_string()), ("sequence-number".to_string(), sequence_number.to_string()), @@ -98,6 +110,7 @@ impl ManifestListWriter { sequence_number, snapshot_id, None, + compression, ) } @@ -108,7 +121,8 @@ impl ManifestListWriter { parent_snapshot_id: Option, sequence_number: i64, first_row_id: Option, // Always None for delete manifests - ) -> Self { + compression: CompressionCodec, + ) -> Result { let mut metadata = HashMap::from_iter([ ("snapshot-id".to_string(), snapshot_id.to_string()), ("sequence-number".to_string(), sequence_number.to_string()), @@ -133,6 +147,7 @@ impl ManifestListWriter { sequence_number, snapshot_id, first_row_id, + compression, ) } @@ -143,26 +158,28 @@ impl ManifestListWriter { sequence_number: i64, snapshot_id: i64, first_row_id: Option, - ) -> Self { + compression: CompressionCodec, + ) -> Result { + let codec = to_avro_codec(compression)?; let avro_schema = match format_version { FormatVersion::V1 => &MANIFEST_LIST_AVRO_SCHEMA_V1, FormatVersion::V2 => &MANIFEST_LIST_AVRO_SCHEMA_V2, FormatVersion::V3 => &MANIFEST_LIST_AVRO_SCHEMA_V3, }; - let mut avro_writer = Writer::new(avro_schema, Vec::new()); + let mut avro_writer = Writer::with_codec(avro_schema, Vec::new(), codec); for (key, value) in metadata { avro_writer .add_user_metadata(key, value) .expect("Avro metadata should be added to the writer before the first record."); } - Self { + Ok(Self { format_version, writer, avro_writer, sequence_number, snapshot_id, next_row_id: first_row_id, - } + }) } /// Append manifests to be written. @@ -228,7 +245,8 @@ impl ManifestListWriter { Ok(()) } - /// Returns number of newly assigned first-row-ids, if any. + /// Assigns `manifest.first_row_id` if not already set, advancing `self.next_row_id` + /// by the manifest's record count. fn assign_first_row_id(&mut self, manifest: &mut ManifestFile) -> Result<()> { match manifest.content { ManifestContentType::Data => { @@ -304,11 +322,12 @@ mod test { use tempfile::TempDir; use super::ManifestListWriter; + use crate::compression::CompressionCodec; use crate::encryption::kms::{KeyManagementClient, MemoryKeyManagementClient}; use crate::encryption::{EncryptedInputFile, EncryptionManager}; use crate::io::{FileIO, FileWrite}; use crate::spec::{ - Datum, FieldSummary, ManifestContentType, ManifestFile, ManifestList, + Datum, FieldSummary, FormatVersion, ManifestContentType, ManifestFile, ManifestList, UNASSIGNED_SEQUENCE_NUMBER, }; @@ -342,7 +361,13 @@ mod test { let io = FileIO::new_with_fs(); let file_writer = file_writer(&path, io).await; - let mut writer = ManifestListWriter::v1(file_writer, 1646658105718557341, Some(0)); + let mut writer = ManifestListWriter::v1( + file_writer, + 1646658105718557341, + Some(0), + CompressionCodec::None, + ) + .unwrap(); writer .add_manifests(expected_manifest_list.entries.clone().into_iter()) .unwrap(); @@ -350,8 +375,7 @@ mod test { let bs = fs::read(path).unwrap(); - let manifest_list = - ManifestList::parse_with_version(&bs, crate::spec::FormatVersion::V1).unwrap(); + let manifest_list = ManifestList::parse_with_version(&bs, FormatVersion::V1).unwrap(); assert_eq!(manifest_list, expected_manifest_list); temp_dir.close().unwrap(); @@ -389,15 +413,21 @@ mod test { let io = FileIO::new_with_fs(); let file_writer = file_writer(&path, io).await; - let mut writer = ManifestListWriter::v2(file_writer, snapshot_id, Some(0), seq_num); + let mut writer = ManifestListWriter::v2( + file_writer, + snapshot_id, + Some(0), + seq_num, + CompressionCodec::None, + ) + .unwrap(); writer .add_manifests(expected_manifest_list.entries.clone().into_iter()) .unwrap(); writer.close().await.unwrap(); let bs = fs::read(path).unwrap(); - let manifest_list = - ManifestList::parse_with_version(&bs, crate::spec::FormatVersion::V2).unwrap(); + let manifest_list = ManifestList::parse_with_version(&bs, FormatVersion::V2).unwrap(); expected_manifest_list.entries[0].sequence_number = seq_num; expected_manifest_list.entries[0].min_sequence_number = seq_num; assert_eq!(manifest_list, expected_manifest_list); @@ -437,16 +467,22 @@ mod test { let io = FileIO::new_with_fs(); let file_writer = file_writer(&path, io).await; - let mut writer = - ManifestListWriter::v3(file_writer, snapshot_id, Some(0), seq_num, Some(10)); + let mut writer = ManifestListWriter::v3( + file_writer, + snapshot_id, + Some(0), + seq_num, + Some(10), + CompressionCodec::None, + ) + .unwrap(); writer .add_manifests(expected_manifest_list.entries.clone().into_iter()) .unwrap(); writer.close().await.unwrap(); let bs = fs::read(path).unwrap(); - let manifest_list = - ManifestList::parse_with_version(&bs, crate::spec::FormatVersion::V3).unwrap(); + let manifest_list = ManifestList::parse_with_version(&bs, FormatVersion::V3).unwrap(); expected_manifest_list.entries[0].sequence_number = seq_num; expected_manifest_list.entries[0].min_sequence_number = seq_num; expected_manifest_list.entries[0].first_row_id = Some(10); @@ -485,7 +521,13 @@ mod test { let io = FileIO::new_with_fs(); let file_writer = file_writer(&path, io).await; - let mut writer = ManifestListWriter::v1(file_writer, 1646658105718557341, Some(0)); + let mut writer = ManifestListWriter::v1( + file_writer, + 1646658105718557341, + Some(0), + CompressionCodec::None, + ) + .unwrap(); writer .add_manifests(expected_manifest_list.entries.clone().into_iter()) .unwrap(); @@ -493,8 +535,7 @@ mod test { let bs = fs::read(path).unwrap(); - let manifest_list = - ManifestList::parse_with_version(&bs, crate::spec::FormatVersion::V2).unwrap(); + let manifest_list = ManifestList::parse_with_version(&bs, FormatVersion::V2).unwrap(); assert_eq!(manifest_list, expected_manifest_list); temp_dir.close().unwrap(); @@ -530,7 +571,13 @@ mod test { let io = FileIO::new_with_fs(); let file_writer = file_writer(&path, io).await; - let mut writer = ManifestListWriter::v1(file_writer, 1646658105718557341, Some(0)); + let mut writer = ManifestListWriter::v1( + file_writer, + 1646658105718557341, + Some(0), + CompressionCodec::None, + ) + .unwrap(); writer .add_manifests(expected_manifest_list.entries.clone().into_iter()) .unwrap(); @@ -538,8 +585,7 @@ mod test { let bs = fs::read(path).unwrap(); - let manifest_list = - ManifestList::parse_with_version(&bs, crate::spec::FormatVersion::V3).unwrap(); + let manifest_list = ManifestList::parse_with_version(&bs, FormatVersion::V3).unwrap(); assert_eq!(manifest_list, expected_manifest_list); temp_dir.close().unwrap(); @@ -577,7 +623,14 @@ mod test { let io = FileIO::new_with_fs(); let file_writer = file_writer(&path, io).await; - let mut writer = ManifestListWriter::v2(file_writer, snapshot_id, Some(0), seq_num); + let mut writer = ManifestListWriter::v2( + file_writer, + snapshot_id, + Some(0), + seq_num, + CompressionCodec::None, + ) + .unwrap(); writer .add_manifests(expected_manifest_list.entries.clone().into_iter()) .unwrap(); @@ -585,8 +638,7 @@ mod test { let bs = fs::read(path).unwrap(); - let manifest_list = - ManifestList::parse_with_version(&bs, crate::spec::FormatVersion::V3).unwrap(); + let manifest_list = ManifestList::parse_with_version(&bs, FormatVersion::V3).unwrap(); expected_manifest_list.entries[0].sequence_number = seq_num; expected_manifest_list.entries[0].min_sequence_number = seq_num; assert_eq!(manifest_list, expected_manifest_list); @@ -630,7 +682,15 @@ mod test { }; let file_writer = encrypted_output.writer().await.unwrap(); - let mut writer = ManifestListWriter::v3(file_writer, snapshot_id, Some(0), seq_num, None); + let mut writer = ManifestListWriter::v3( + file_writer, + snapshot_id, + Some(0), + seq_num, + None, + CompressionCodec::None, + ) + .unwrap(); writer .add_manifests(expected.entries.clone().into_iter()) .unwrap(); @@ -638,7 +698,7 @@ mod test { let raw_bytes = file_io.new_input(path).unwrap().read().await.unwrap(); assert!( - ManifestList::parse_with_version(&raw_bytes, crate::spec::FormatVersion::V3).is_err(), + ManifestList::parse_with_version(&raw_bytes, FormatVersion::V3).is_err(), "raw bytes should be ciphertext, not parseable as Avro" ); @@ -649,7 +709,7 @@ mod test { .await .unwrap(); let manifest_list = - ManifestList::parse_with_version(&plaintext, crate::spec::FormatVersion::V3).unwrap(); + ManifestList::parse_with_version(&plaintext, FormatVersion::V3).unwrap(); expected.entries[0].sequence_number = seq_num; expected.entries[0].min_sequence_number = seq_num; @@ -666,6 +726,93 @@ mod test { (mgr, FileIO::new_with_memory()) } + #[tokio::test] + async fn test_manifest_list_writer_with_compression() { + let snapshot_id = 377075049360453639; + let seq_num = 1; + + let entries: Vec = (0..1000) + .map(|i| ManifestFile { + manifest_path: format!( + "s3a://icebergdata/demo/s1/t1/metadata/very-long-path-for-compression-test/manifest-file-number-{i}.avro" + ), + manifest_length: 6926 + i, + partition_spec_id: 1, + content: ManifestContentType::Data, + sequence_number: seq_num, + min_sequence_number: seq_num, + added_snapshot_id: snapshot_id, + added_files_count: Some(1), + existing_files_count: Some(0), + deleted_files_count: Some(0), + added_rows_count: Some(3), + existing_rows_count: Some(0), + deleted_rows_count: Some(0), + partitions: None, + key_metadata: None, + first_row_id: None, + }) + .collect(); + + let tmp_dir = TempDir::new().unwrap(); + let io = FileIO::new_with_fs(); + + // Write the same entries with a given codec and return (file size, round-tripped entry count). + async fn write_and_read( + io: &FileIO, + path: &Path, + snapshot_id: i64, + seq_num: i64, + entries: &[ManifestFile], + compression: CompressionCodec, + ) -> (u64, usize) { + let mut writer = ManifestListWriter::v2( + file_writer(path, io.clone()).await, + snapshot_id, + Some(0), + seq_num, + compression, + ) + .unwrap(); + writer.add_manifests(entries.iter().cloned()).unwrap(); + writer.close().await.unwrap(); + + let size = fs::metadata(path).unwrap().len(); + let bytes = fs::read(path).unwrap(); + let manifest_list = + ManifestList::parse_with_version(&bytes, FormatVersion::V2).unwrap(); + (size, manifest_list.entries().len()) + } + + let (uncompressed_size, uncompressed_count) = write_and_read( + &io, + &tmp_dir.path().join("manifest_list_uncompressed.avro"), + snapshot_id, + seq_num, + &entries, + CompressionCodec::None, + ) + .await; + assert_eq!(uncompressed_count, 1000); + + // Every compressed codec must shrink the file and still round-trip all entries. + for codec in [CompressionCodec::Gzip(9), CompressionCodec::Zstd(3)] { + let path = tmp_dir + .path() + .join(format!("manifest_list_{}.avro", codec.name())); + let (compressed_size, compressed_count) = + write_and_read(&io, &path, snapshot_id, seq_num, &entries, codec).await; + assert_eq!( + compressed_count, 1000, + "{codec} should round-trip all entries" + ); + assert!( + compressed_size < uncompressed_size, + "{codec} size ({compressed_size}) should be less than uncompressed size ({uncompressed_size})" + ); + } + } + async fn file_writer(path: &Path, io: FileIO) -> Box { io.new_output(path.to_str().unwrap()) .unwrap() diff --git a/crates/iceberg/src/spec/mod.rs b/crates/iceberg/src/spec/mod.rs index 707ebbb630..fbae9c6c3b 100644 --- a/crates/iceberg/src/spec/mod.rs +++ b/crates/iceberg/src/spec/mod.rs @@ -17,6 +17,7 @@ //! Spec for Iceberg. +pub(crate) mod avro_util; mod datatypes; mod encrypted_key; mod manifest; diff --git a/crates/iceberg/src/spec/table_properties.rs b/crates/iceberg/src/spec/table_properties.rs index f418d3a1de..8393921de8 100644 --- a/crates/iceberg/src/spec/table_properties.rs +++ b/crates/iceberg/src/spec/table_properties.rs @@ -22,7 +22,7 @@ use iceberg_property_macro::properties_view; use crate::compression::CompressionCodec; use crate::encryption::AesKeySize; use crate::error::{Result, invalid_data}; -use crate::spec::NameMapping; +use crate::spec::{NameMapping, avro_util}; use crate::util::location::strip_trailing_slash; fn parse_location_property(path: &str) -> Result { @@ -86,21 +86,46 @@ fn parse_parquet_compression( .transpose()? .unwrap_or(default); + apply_compression_level(properties, level_key, codec) +} + +/// Parse the Avro compression codec for manifests and manifest lists +/// (`write.avro.compression-codec`) and fold in the compression level +/// (`write.avro.compression-level`) for the codecs that accept one (`zstd`, `gzip`). +fn parse_avro_compression( + properties: &HashMap, + codec_key: &str, + additional_keys: &[&str], + default: CompressionCodec, +) -> Result { + let codec = properties + .get(codec_key) + .map(|value| avro_util::parse_avro_codec(value)) + .transpose()? + .unwrap_or(default); + let codec = apply_compression_level(properties, additional_keys[0], codec)?; + // Reject codecs and levels that Avro cannot write (e.g. lz4, gzip level 10) when the + // property is read, rather than deep inside a later commit. + avro_util::to_avro_codec(codec)?; + Ok(codec) +} + +/// Override the level of `codec` with the value of `level_key`, if set. +fn apply_compression_level( + properties: &HashMap, + level_key: &str, + codec: CompressionCodec, +) -> Result { let level = properties .get(level_key) .map(|value| { value .parse::() - .map_err(|error| invalid_data!("Invalid value for {level_key}: {error}")) + .map_err(|error| invalid_data!("Invalid value for {level_key}: {value} ({error})")) }) .transpose()?; - Ok(match (codec, level) { - (CompressionCodec::Zstd(_), Some(level)) => CompressionCodec::Zstd(level), - (CompressionCodec::Gzip(_), Some(level)) => CompressionCodec::Gzip(level), - (CompressionCodec::Brotli(_), Some(level)) => CompressionCodec::Brotli(level), - (codec, _) => codec, - }) + Ok(level.map_or(codec, |level| codec.with_level(level))) } properties_view! { @@ -166,6 +191,17 @@ pub struct TableProperties { getter )] metadata_compression_codec: CompressionCodec, + /// Compression codec for Avro files (manifests, manifest lists), with the + /// resolved compression level folded in (from `write.avro.compression-level`, + /// or the codec's Avro default when unset). + #[property( + key = Self::PROPERTY_AVRO_COMPRESSION_CODEC, + additional_keys = [Self::PROPERTY_AVRO_COMPRESSION_LEVEL], + default = avro_util::DEFAULT_AVRO_CODEC, + parse_properties_with = parse_avro_compression, + getter + )] + avro_compression_codec: CompressionCodec, /// Whether to use `FanoutWriter` for partitioned tables. #[property( key = Self::PROPERTY_DATAFUSION_WRITE_FANOUT_ENABLED, @@ -444,6 +480,14 @@ impl TableProperties<'_> { "write.metadata.compression-codec"; /// Default metadata compression codec - uncompressed pub const PROPERTY_METADATA_COMPRESSION_CODEC_DEFAULT: &'static str = "none"; + + /// Compression codec for Avro files (manifests, manifest lists) + pub const PROPERTY_AVRO_COMPRESSION_CODEC: &'static str = "write.avro.compression-codec"; + /// Default Avro compression codec - gzip + pub const PROPERTY_AVRO_COMPRESSION_CODEC_DEFAULT: &'static str = "gzip"; + /// Compression level for Avro files, interpreted per the configured codec (gzip 0–9, + /// zstd 0–22). Ignored by codecs that carry no level. + pub const PROPERTY_AVRO_COMPRESSION_LEVEL: &'static str = "write.avro.compression-level"; /// Whether to use `FanoutWriter` for partitioned tables (handles unsorted data). /// If false, uses `ClusteredWriter` (requires sorted data, more memory efficient). pub const PROPERTY_DATAFUSION_WRITE_FANOUT_ENABLED: &'static str = @@ -590,6 +634,12 @@ mod tests { table_properties.metadata_compression_codec().unwrap(), CompressionCodec::None ); + // Keeps the public default codec name in sync with `avro_util::DEFAULT_AVRO_CODEC`. + assert_eq!( + table_properties.avro_compression_codec().unwrap(), + avro_util::parse_avro_codec(TableProperties::PROPERTY_AVRO_COMPRESSION_CODEC_DEFAULT) + .unwrap() + ); assert_eq!( table_properties.gc_enabled().unwrap(), TableProperties::PROPERTY_GC_ENABLED_DEFAULT @@ -698,6 +748,69 @@ mod tests { } } + fn avro_compression_props(codec: Option<&str>, level: Option<&str>) -> HashMap { + [ + codec.map(|c| (TableProperties::PROPERTY_AVRO_COMPRESSION_CODEC, c)), + level.map(|l| (TableProperties::PROPERTY_AVRO_COMPRESSION_LEVEL, l)), + ] + .into_iter() + .flatten() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect() + } + + #[test] + fn test_table_properties_avro_compression() { + for (codec, level, expected) in [ + (Some("zstd"), Some("3"), CompressionCodec::Zstd(3)), + ( + Some("zstd"), + None, + avro_util::parse_avro_codec("zstd").unwrap(), + ), + // A level without a codec applies to the default (gzip) codec. + (None, Some("5"), CompressionCodec::Gzip(5)), + // Codecs without a level ignore it. + (Some("snappy"), Some("5"), CompressionCodec::Snappy), + // zstd levels above 22 are accepted; the writer clamps them, as zstd does. + (Some("zstd"), Some("200"), CompressionCodec::Zstd(200)), + ] { + let props = avro_compression_props(codec, level); + assert_eq!( + TableProperties::new(&props) + .avro_compression_codec() + .unwrap(), + expected, + "codec {codec:?}, level {level:?}" + ); + } + } + + #[test] + fn test_table_properties_avro_compression_invalid() { + // Each case must fail when the property is read, with an error naming the bad value. + for (codec, level, bad_value) in [ + // Not a number. + (None, Some("abc"), "abc"), + // Overflows u8. + (None, Some("300"), "300"), + // Above Java's Deflater range of 0-9. + (Some("gzip"), Some("10"), "10"), + // No Avro representation. + (Some("lz4"), None, "lz4"), + ] { + let props = avro_compression_props(codec, level); + let err = TableProperties::new(&props) + .avro_compression_codec() + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + assert!( + err.to_string().contains(bad_value), + "unexpected error for codec {codec:?}, level {level:?}: {err}" + ); + } + } + #[test] fn test_table_properties_compression() { let props = HashMap::from([( diff --git a/crates/iceberg/src/test_utils/scan.rs b/crates/iceberg/src/test_utils/scan.rs index 0c4a7e9dd7..0fb99067dd 100644 --- a/crates/iceberg/src/test_utils/scan.rs +++ b/crates/iceberg/src/test_utils/scan.rs @@ -35,6 +35,7 @@ use tempfile::TempDir; use uuid::Uuid; use crate::TableIdent; +use crate::compression::CompressionCodec; use crate::io::{FileIO, OutputFile}; use crate::metadata_columns::{ RESERVED_COL_NAME_DELETE_FILE_PATH, RESERVED_COL_NAME_DELETE_FILE_POS, @@ -294,7 +295,9 @@ impl TableTestFixture { Some(current_snapshot.snapshot_id()), current_schema.clone(), current_partition_spec.as_ref().clone(), + CompressionCodec::None, ) + .unwrap() .build_v2_data(); writer .add_entry( @@ -376,7 +379,9 @@ impl TableTestFixture { current_snapshot.snapshot_id(), current_snapshot.parent_snapshot_id(), current_snapshot.sequence_number(), - ); + CompressionCodec::None, + ) + .unwrap(); manifest_list_write .add_manifests(vec![data_file_manifest].into_iter()) .unwrap(); @@ -416,7 +421,9 @@ impl TableTestFixture { Some(current_snapshot.snapshot_id()), current_schema.clone(), current_partition_spec.as_ref().clone(), + CompressionCodec::None, ) + .unwrap() .build_v3_data(); writer .add_entry( @@ -454,7 +461,9 @@ impl TableTestFixture { current_snapshot.parent_snapshot_id(), current_snapshot.sequence_number(), Some(42), - ); + CompressionCodec::None, + ) + .unwrap(); manifest_list_write .add_manifests(vec![data_file_manifest].into_iter()) .unwrap(); @@ -477,7 +486,9 @@ impl TableTestFixture { Some(current_snapshot.snapshot_id()), current_schema.clone(), current_partition_spec.as_ref().clone(), + CompressionCodec::None, ) + .unwrap() .build_v2_data(); writer .add_entry( @@ -571,7 +582,9 @@ impl TableTestFixture { current_snapshot.snapshot_id(), current_snapshot.parent_snapshot_id(), current_snapshot.sequence_number(), - ); + CompressionCodec::None, + ) + .unwrap(); manifest_list_write .add_manifests(vec![data_file_manifest].into_iter()) .unwrap(); @@ -709,7 +722,9 @@ impl TableTestFixture { Some(current_snapshot.snapshot_id()), current_schema.clone(), current_partition_spec.as_ref().clone(), + CompressionCodec::None, ) + .unwrap() .build_v2_data(); // Create an empty partition value. @@ -798,7 +813,9 @@ impl TableTestFixture { current_snapshot.snapshot_id(), current_snapshot.parent_snapshot_id(), current_snapshot.sequence_number(), - ); + CompressionCodec::None, + ) + .unwrap(); manifest_list_write .add_manifests(vec![data_file_manifest].into_iter()) .unwrap(); @@ -819,7 +836,9 @@ impl TableTestFixture { Some(current_snapshot.snapshot_id()), current_schema.clone(), current_partition_spec.as_ref().clone(), + CompressionCodec::None, ) + .unwrap() .build_v2_data(); // Add 10 data entries @@ -853,7 +872,9 @@ impl TableTestFixture { Some(current_snapshot.snapshot_id()), current_schema.clone(), current_partition_spec.as_ref().clone(), + CompressionCodec::None, ) + .unwrap() .build_v2_deletes(); writer @@ -892,7 +913,9 @@ impl TableTestFixture { current_snapshot.snapshot_id(), current_snapshot.parent_snapshot_id(), current_snapshot.sequence_number(), - ); + CompressionCodec::None, + ) + .unwrap(); manifest_list_write .add_manifests(vec![data_manifest, delete_manifest].into_iter()) .unwrap(); @@ -925,7 +948,9 @@ impl TableTestFixture { Some(current_snapshot.snapshot_id()), current_schema.clone(), current_partition_spec.as_ref().clone(), + CompressionCodec::None, ) + .unwrap() .build_v2_data(); data_writer .add_entry( @@ -960,7 +985,9 @@ impl TableTestFixture { Some(current_snapshot.snapshot_id()), current_schema.clone(), current_partition_spec.as_ref().clone(), + CompressionCodec::None, ) + .unwrap() .build_v2_deletes(); delete_writer .add_entry( @@ -997,7 +1024,9 @@ impl TableTestFixture { current_snapshot.snapshot_id(), current_snapshot.parent_snapshot_id(), current_snapshot.sequence_number(), - ); + CompressionCodec::None, + ) + .unwrap(); manifest_list_write .add_manifests(manifests.into_iter()) .unwrap(); @@ -1022,7 +1051,9 @@ impl TableTestFixture { Some(current_snapshot.snapshot_id()), current_schema.clone(), current_partition_spec.as_ref().clone(), + CompressionCodec::None, ) + .unwrap() .build_v2_data(); for (i, sort_order_id) in sort_order_ids.into_iter().enumerate() { @@ -1067,7 +1098,9 @@ impl TableTestFixture { current_snapshot.snapshot_id(), current_snapshot.parent_snapshot_id(), current_snapshot.sequence_number(), - ); + CompressionCodec::None, + ) + .unwrap(); manifest_list_write .add_manifests(std::iter::once(data_file_manifest)) .unwrap(); diff --git a/crates/iceberg/src/transaction/append.rs b/crates/iceberg/src/transaction/append.rs index 3260f9ffa2..e3efffde83 100644 --- a/crates/iceberg/src/transaction/append.rs +++ b/crates/iceberg/src/transaction/append.rs @@ -161,6 +161,7 @@ mod tests { use tempfile::TempDir; use uuid::Uuid; + use crate::compression::CompressionCodec; use crate::encryption::kms::MemoryKeyManagementClient; use crate::encryption::{SensitiveBytes, StandardKeyMetadata}; use crate::io::FileIO; @@ -240,7 +241,9 @@ mod tests { Some(current_snapshot.snapshot_id()), schema.clone(), partition_spec.as_ref().clone(), + CompressionCodec::None, ) + .unwrap() .build_v2_data(); data_writer .add_entry( @@ -269,7 +272,9 @@ mod tests { Some(current_snapshot.snapshot_id()), schema.clone(), partition_spec.as_ref().clone(), + CompressionCodec::None, ) + .unwrap() .build_v2_data(); delete_writer .add_delete_entry( @@ -311,7 +316,9 @@ mod tests { current_snapshot.snapshot_id(), current_snapshot.parent_snapshot_id(), current_snapshot.sequence_number(), - ); + CompressionCodec::None, + ) + .unwrap(); manifest_list_writer .add_manifests(vec![data_manifest, delete_manifest].into_iter()) .unwrap(); diff --git a/crates/iceberg/src/transaction/snapshot.rs b/crates/iceberg/src/transaction/snapshot.rs index a0ce1fba14..5954199c93 100644 --- a/crates/iceberg/src/transaction/snapshot.rs +++ b/crates/iceberg/src/transaction/snapshot.rs @@ -23,6 +23,7 @@ use futures::TryStreamExt; use futures::stream::FuturesUnordered; use uuid::Uuid; +use crate::compression::CompressionCodec; use crate::error::{Result, invalid_data}; use crate::spec::{ DataFile, DataFileFormat, FormatVersion, MAIN_BRANCH, ManifestContentType, ManifestEntry, @@ -158,6 +159,13 @@ impl<'a> SnapshotProducer<'a> { Ok(()) } + fn avro_compression_codec(&self) -> Result { + self.table + .metadata() + .table_properties() + .avro_compression_codec() + } + pub(crate) async fn validate_duplicate_files(&self) -> Result<()> { let Some(current_snapshot) = self.table.metadata().current_snapshot() else { return Ok(()); @@ -245,6 +253,7 @@ impl<'a> SnapshotProducer<'a> { .as_ref() .clone(); let schema = self.table.metadata().current_schema().clone(); + let compression = self.avro_compression_codec()?; let builder = if let Some(em) = self.table.encryption_manager() { ManifestWriterBuilder::new_from_encrypted( @@ -252,9 +261,16 @@ impl<'a> SnapshotProducer<'a> { Some(self.snapshot_id), schema, partition_spec, + compression, )? } else { - ManifestWriterBuilder::new(output_file, Some(self.snapshot_id), schema, partition_spec) + ManifestWriterBuilder::new( + output_file, + Some(self.snapshot_id), + schema, + partition_spec, + compression, + )? }; match self.table.metadata().format_version() { @@ -456,20 +472,27 @@ impl<'a> SnapshotProducer<'a> { }; let parent_snapshot_id = self.table.metadata().current_snapshot_id(); + let compression = self.avro_compression_codec()?; + let mut manifest_list_writer = match self.table.metadata().format_version() { FormatVersion::V1 => { - ManifestListWriter::v1(writer, self.snapshot_id, parent_snapshot_id) - } - FormatVersion::V2 => { - ManifestListWriter::v2(writer, self.snapshot_id, parent_snapshot_id, next_seq_num) + ManifestListWriter::v1(writer, self.snapshot_id, parent_snapshot_id, compression)? } + FormatVersion::V2 => ManifestListWriter::v2( + writer, + self.snapshot_id, + parent_snapshot_id, + next_seq_num, + compression, + )?, FormatVersion::V3 => ManifestListWriter::v3( writer, self.snapshot_id, parent_snapshot_id, next_seq_num, Some(first_row_id), - ), + compression, + )?, }; // Calling self.summary() before self.produce_manifests() is important because self.added_data_files