From a75b1f97166ce531f3fc504b2cd933031e1bc447 Mon Sep 17 00:00:00 2001 From: xuanyili Date: Sat, 25 Jul 2026 23:05:59 +0000 Subject: [PATCH 1/2] feat: support zstd metadata compression Add zstd metadata codec validation, canonical filename generation, alternate suffix parsing, and content-based frame detection for reads. Cover interoperability, codec transitions, truncated frames, and semantic round trips. Part of apache/iceberg-rust#2411. --- .../iceberg/src/catalog/metadata_location.rs | 191 +++++++++++++----- crates/iceberg/src/compression.rs | 14 +- crates/iceberg/src/spec/table_metadata.rs | 155 ++++++++++++-- crates/iceberg/src/spec/table_properties.rs | 170 ++++++---------- .../TableMetadataV2Valid.zstd.metadata.json | Bin 0 -> 374 bytes 5 files changed, 342 insertions(+), 188 deletions(-) create mode 100644 crates/iceberg/testdata/table_metadata/TableMetadataV2Valid.zstd.metadata.json diff --git a/crates/iceberg/src/catalog/metadata_location.rs b/crates/iceberg/src/catalog/metadata_location.rs index d5daa02b37..025fbda426 100644 --- a/crates/iceberg/src/catalog/metadata_location.rs +++ b/crates/iceberg/src/catalog/metadata_location.rs @@ -30,7 +30,7 @@ use crate::{Error, ErrorKind, Result}; pub(crate) const METADATA_FOLDER_NAME: &str = "metadata"; /// Helper for parsing a location of the format: `/-.metadata.json` -/// or with compression: `/-.gz.metadata.json` +/// or with compression: `/-..metadata.json` /// /// `` is set to the `write.metadata.path` table property and /// it defaults to the `/metadata` when the property is not set. @@ -90,20 +90,34 @@ impl MetadataLocation { } /// Parses a file name of the format `-.metadata.json` - /// or with compression: `-.gz.metadata.json`. + /// or with compression before or after `.metadata.json`. /// Parse errors for compression codec result in CompressionCodec::None. fn parse_file_name(file_name: &str) -> Result<(i32, Uuid, CompressionCodec)> { - let stripped = file_name.strip_suffix(".metadata.json").ok_or(Error::new( - ErrorKind::Unexpected, - format!("Invalid metadata file ending: {file_name}"), - ))?; - - // Check for compression suffix (e.g., .gz) let gzip_suffix = CompressionCodec::gzip_default().suffix()?; - let (stripped, compression_codec) = if let Some(s) = stripped.strip_suffix(gzip_suffix) { - (s, CompressionCodec::gzip_default()) + let zstd_suffix = CompressionCodec::zstd_default().suffix()?; + let metadata_suffix = ".metadata.json"; + + let (stripped, compression_codec) = if let Some(stripped) = + file_name.strip_suffix(&format!("{metadata_suffix}{zstd_suffix}")) + { + (stripped, CompressionCodec::zstd_default()) + } else if let Some(stripped) = + file_name.strip_suffix(&format!("{metadata_suffix}{gzip_suffix}")) + { + (stripped, CompressionCodec::gzip_default()) + } else if let Some(stripped) = file_name.strip_suffix(metadata_suffix) { + if let Some(stripped) = stripped.strip_suffix(zstd_suffix) { + (stripped, CompressionCodec::zstd_default()) + } else if let Some(stripped) = stripped.strip_suffix(gzip_suffix) { + (stripped, CompressionCodec::gzip_default()) + } else { + (stripped, CompressionCodec::None) + } } else { - (stripped, CompressionCodec::None) + return Err(Error::new( + ErrorKind::Unexpected, + format!("Invalid metadata file ending: {file_name}"), + )); }; let (version, id) = stripped.split_once('-').ok_or(Error::new( @@ -248,6 +262,36 @@ mod test { compression_codec: CompressionCodec::gzip_default(), }), ), + // With trailing gzip compression suffix + ( + "/abc/metadata/1234567-2cd22b57-5127-4198-92ba-e4e67c79821b.metadata.json.gz", + Ok(MetadataLocation { + location: "/abc/metadata".to_string(), + version: 1234567, + id: Uuid::from_str("2cd22b57-5127-4198-92ba-e4e67c79821b").unwrap(), + compression_codec: CompressionCodec::gzip_default(), + }), + ), + // With zstd compression + ( + "/abc/metadata/1234567-2cd22b57-5127-4198-92ba-e4e67c79821b.zstd.metadata.json", + Ok(MetadataLocation { + location: "/abc/metadata".to_string(), + version: 1234567, + id: Uuid::from_str("2cd22b57-5127-4198-92ba-e4e67c79821b").unwrap(), + compression_codec: CompressionCodec::zstd_default(), + }), + ), + // With trailing zstd compression suffix + ( + "/abc/metadata/1234567-2cd22b57-5127-4198-92ba-e4e67c79821b.metadata.json.zstd", + Ok(MetadataLocation { + location: "/abc/metadata".to_string(), + version: 1234567, + id: Uuid::from_str("2cd22b57-5127-4198-92ba-e4e67c79821b").unwrap(), + compression_codec: CompressionCodec::zstd_default(), + }), + ), // Negative version ( "/metadata/-123-2cd22b57-5127-4198-92ba-e4e67c79821b.metadata.json", @@ -295,6 +339,25 @@ mod test { } } + #[test] + fn test_metadata_location_canonicalizes_compression_suffixes() { + for (input, expected) in [ + ( + "/abc/metadata/00001-2cd22b57-5127-4198-92ba-e4e67c79821b.metadata.json.gz", + "/abc/metadata/00001-2cd22b57-5127-4198-92ba-e4e67c79821b.gz.metadata.json", + ), + ( + "/abc/metadata/00001-2cd22b57-5127-4198-92ba-e4e67c79821b.metadata.json.zstd", + "/abc/metadata/00001-2cd22b57-5127-4198-92ba-e4e67c79821b.zstd.metadata.json", + ), + ] { + assert_eq!( + MetadataLocation::from_str(input).unwrap().to_string(), + expected + ); + } + } + #[test] fn test_metadata_location_with_next_version() { let metadata = create_test_metadata(HashMap::new()); @@ -318,33 +381,34 @@ mod test { #[test] fn test_with_next_version_preserves_compression() { - // Start from a parsed location with no compression - let location_none = MetadataLocation::from_str( - "/test/table/metadata/00000-2cd22b57-5127-4198-92ba-e4e67c79821b.metadata.json", - ) - .unwrap(); - assert_eq!(location_none.compression_codec, CompressionCodec::None); - - let next_none = location_none.with_next_version(); - assert_eq!(next_none.compression_codec, CompressionCodec::None); - assert_eq!(next_none.version, 1); - - // Start from a parsed location with gzip compression - let location_gzip = MetadataLocation::from_str( - "/test/table/metadata/00005-81056704-ce5b-41c4-bb83-eb6408081af6.gz.metadata.json", - ) - .unwrap(); - assert_eq!( - location_gzip.compression_codec, - CompressionCodec::gzip_default() - ); - - let next_gzip = location_gzip.with_next_version(); - assert_eq!( - next_gzip.compression_codec, - CompressionCodec::gzip_default() - ); - assert_eq!(next_gzip.version, 6); + for (input, expected_codec, expected_version, expected_suffix) in [ + ( + "/test/table/metadata/00000-2cd22b57-5127-4198-92ba-e4e67c79821b.metadata.json", + CompressionCodec::None, + 1, + ".metadata.json", + ), + ( + "/test/table/metadata/00005-81056704-ce5b-41c4-bb83-eb6408081af6.metadata.json.gz", + CompressionCodec::gzip_default(), + 6, + ".gz.metadata.json", + ), + ( + "/test/table/metadata/00009-81056704-ce5b-41c4-bb83-eb6408081af6.metadata.json.zstd", + CompressionCodec::zstd_default(), + 10, + ".zstd.metadata.json", + ), + ] { + let location = MetadataLocation::from_str(input).unwrap(); + assert_eq!(location.compression_codec, expected_codec); + + let next = location.with_next_version(); + assert_eq!(next.compression_codec, expected_codec); + assert_eq!(next.version, expected_version); + assert!(next.to_string().ends_with(expected_suffix)); + } } #[test] @@ -374,10 +438,26 @@ mod test { "/test/table/metadata/00000-2cd22b57-5127-4198-92ba-e4e67c79821b.gz.metadata.json" ); + // Transition from gzip to zstd compression + let props_zstd = HashMap::from([( + TableProperties::PROPERTY_METADATA_COMPRESSION_CODEC.to_string(), + "zstd".to_string(), + )]); + let metadata_zstd = create_test_metadata(props_zstd); + let updated_zstd = updated_gzip.try_with_new_metadata(&metadata_zstd).unwrap(); + assert_eq!( + updated_zstd.compression_codec, + CompressionCodec::zstd_default() + ); + assert_eq!( + updated_zstd.to_string(), + "/test/table/metadata/00000-2cd22b57-5127-4198-92ba-e4e67c79821b.zstd.metadata.json" + ); + // Update back to no compression let props_none = HashMap::new(); let metadata_none = create_test_metadata(props_none); - let updated_none = updated_gzip.try_with_new_metadata(&metadata_none).unwrap(); + let updated_none = updated_zstd.try_with_new_metadata(&metadata_none).unwrap(); assert_eq!(updated_none.compression_codec, CompressionCodec::None); assert_eq!(updated_none.version, 0); assert_eq!( @@ -436,15 +516,30 @@ mod test { #[test] fn test_new_with_metadata_honors_write_metadata_path() { - // Test metadata lives under `/metadata` by default - let default_meta = create_test_metadata(HashMap::new()); - let default_loc = MetadataLocation::try_new_with_metadata(&default_meta).unwrap(); - assert!( - default_loc - .to_string() - .starts_with("/test/table/metadata/00000-"), - "unexpected location: {default_loc}" - ); + // Test metadata lives under `/metadata` with canonical compression suffixes. + for (compression, expected_suffix) in [ + (None, ".metadata.json"), + (Some("gzip"), ".gz.metadata.json"), + (Some("zstd"), ".zstd.metadata.json"), + ] { + let properties = compression + .map(|compression| { + HashMap::from([( + TableProperties::PROPERTY_METADATA_COMPRESSION_CODEC.to_string(), + compression.to_string(), + )]) + }) + .unwrap_or_default(); + let metadata = create_test_metadata(properties); + let location = MetadataLocation::try_new_with_metadata(&metadata).unwrap(); + assert!( + location + .to_string() + .starts_with("/test/table/metadata/00000-"), + "unexpected location: {location}" + ); + assert!(location.to_string().ends_with(expected_suffix)); + } // Test a configured `write.metadata.path` is honored let props = HashMap::from([( diff --git a/crates/iceberg/src/compression.rs b/crates/iceberg/src/compression.rs index 929d9226e7..266d93fbe8 100644 --- a/crates/iceberg/src/compression.rs +++ b/crates/iceberg/src/compression.rs @@ -169,18 +169,17 @@ impl CompressionCodec { } /// Returns the file extension suffix for this compression codec. - /// Returns empty string for None, ".gz" for Gzip. + /// Returns empty string for None, ".zstd" for Zstd, and ".gz" for Gzip. /// /// # Errors /// - /// Returns an error for Lz4 and Zstd as they are not fully supported. + /// Returns an error for Lz4 and Snappy as they are not fully supported. pub fn suffix(&self) -> Result<&'static str> { match self { CompressionCodec::None => Ok(""), + CompressionCodec::Zstd(_) => Ok(".zstd"), CompressionCodec::Gzip(_) => Ok(".gz"), - codec @ (CompressionCodec::Lz4 - | CompressionCodec::Zstd(_) - | CompressionCodec::Snappy) => Err(Error::new( + codec @ (CompressionCodec::Lz4 | CompressionCodec::Snappy) => Err(Error::new( ErrorKind::FeatureUnsupported, format!("suffix not defined for {codec:?}"), )), @@ -244,17 +243,14 @@ mod tests { #[test] fn test_suffix() { assert_eq!(CompressionCodec::None.suffix().unwrap(), ""); + assert_eq!(CompressionCodec::zstd_default().suffix().unwrap(), ".zstd"); 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] diff --git a/crates/iceberg/src/spec/table_metadata.rs b/crates/iceberg/src/spec/table_metadata.rs index ecc0586680..0733bac1f6 100644 --- a/crates/iceberg/src/spec/table_metadata.rs +++ b/crates/iceberg/src/spec/table_metadata.rs @@ -46,6 +46,8 @@ use crate::{Error, ErrorKind}; static MAIN_BRANCH: &str = "main"; pub(crate) static ONE_MINUTE_MS: i64 = 60_000; +const GZIP_MAGIC: &[u8] = &[0x1F, 0x8B]; +const ZSTD_MAGIC: &[u8] = &[0x28, 0xB5, 0x2F, 0xFD]; /// Sentinel value used by the Java implementation and older metadata files /// to represent a missing/empty current snapshot ID. During deserialization, @@ -378,7 +380,7 @@ impl TableMetadata { /// Returns the metadata compression codec from table properties. /// /// Returns `CompressionCodec::None` if compression is disabled or not configured. - /// Returns `CompressionCodec::Gzip` if gzip compression is enabled. + /// Returns the configured gzip or zstd codec when compression is enabled. /// /// # Errors /// @@ -466,21 +468,23 @@ impl TableMetadata { let input_file = file_io.new_input(metadata_location)?; let metadata_content = input_file.read().await?; - // Check if the file is compressed by looking for the gzip "magic number". - let metadata = if metadata_content.len() > 2 - && metadata_content[0] == 0x1F - && metadata_content[1] == 0x8B - { - let decompressed_data = CompressionCodec::gzip_default() - .decompress(metadata_content.to_vec()) - .map_err(|e| { - Error::new( - ErrorKind::DataInvalid, - "Trying to read compressed metadata file", - ) - .with_context("file_path", metadata_location) - .with_source(e) - })?; + let compression_codec = if metadata_content.starts_with(GZIP_MAGIC) { + Some(CompressionCodec::gzip_default()) + } else if metadata_content.starts_with(ZSTD_MAGIC) { + Some(CompressionCodec::zstd_default()) + } else { + None + }; + + let metadata = if let Some(codec) = compression_codec { + let decompressed_data = codec.decompress(metadata_content.to_vec()).map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + "Trying to read compressed metadata file", + ) + .with_context("file_path", metadata_location) + .with_source(e) + })?; serde_json::from_slice(&decompressed_data)? } else { serde_json::from_slice(&metadata_content)? @@ -513,7 +517,7 @@ impl TableMetadata { // Apply compression based on codec let data_to_write = match codec { - CompressionCodec::Gzip(_) => codec.compress(json_data)?, + CompressionCodec::Gzip(_) | CompressionCodec::Zstd(_) => codec.compress(json_data)?, CompressionCodec::None => json_data, _ => { return Err(Error::new( @@ -1629,7 +1633,7 @@ mod tests { use tempfile::TempDir; use uuid::Uuid; - use super::{FormatVersion, MetadataLog, SnapshotLog, TableMetadataBuilder}; + use super::{FormatVersion, MetadataLog, SnapshotLog, TableMetadataBuilder, ZSTD_MAGIC}; use crate::catalog::MetadataLocation; use crate::compression::CompressionCodec; use crate::io::FileIO; @@ -3680,6 +3684,80 @@ mod tests { assert_eq!(read_metadata, original_metadata); } + #[tokio::test] + async fn test_table_metadata_read_iceberg_go_zstd_fixture() { + let fixture = + fs::read("testdata/table_metadata/TableMetadataV2Valid.zstd.metadata.json").unwrap(); + assert!(fixture.starts_with(ZSTD_MAGIC)); + + let temp_dir = TempDir::new().unwrap(); + let file_io = FileIO::new_with_fs(); + let mut expected_metadata = None; + + for file_name in [ + "v1.zstd.metadata.json", + "v1.metadata.json.zstd", + "v1.metadata.json", + ] { + let metadata_location = temp_dir.path().join(file_name); + fs::write(&metadata_location, &fixture).unwrap(); + + let read_metadata = + TableMetadata::read_from(&file_io, metadata_location.to_str().unwrap()) + .await + .unwrap(); + if let Some(expected_metadata) = &expected_metadata { + assert_eq!(&read_metadata, expected_metadata); + } else { + assert_eq!(read_metadata.format_version(), FormatVersion::V2); + assert_eq!( + read_metadata.uuid(), + Uuid::parse_str("9c12d441-03fe-4693-9a96-a0705ddf69c1").unwrap() + ); + assert_eq!(read_metadata.location(), "s3://bucket/test/location"); + expected_metadata = Some(read_metadata); + } + } + } + + #[tokio::test] + async fn test_table_metadata_read_plain_json_with_zstd_file_name() { + let temp_dir = TempDir::new().unwrap(); + let metadata_location = temp_dir.path().join("v1.zstd.metadata.json"); + let expected_metadata: TableMetadata = get_test_table_metadata("TableMetadataV2Valid.json"); + fs::write( + &metadata_location, + serde_json::to_vec(&expected_metadata).unwrap(), + ) + .unwrap(); + + let file_io = FileIO::new_with_fs(); + let read_metadata = TableMetadata::read_from(&file_io, metadata_location.to_str().unwrap()) + .await + .unwrap(); + assert_eq!(read_metadata, expected_metadata); + } + + #[tokio::test] + async fn test_table_metadata_read_truncated_zstd() { + let mut fixture = + fs::read("testdata/table_metadata/TableMetadataV2Valid.zstd.metadata.json").unwrap(); + fixture.truncate(fixture.len() / 2); + assert!(fixture.starts_with(ZSTD_MAGIC)); + + let temp_dir = TempDir::new().unwrap(); + let metadata_location = temp_dir.path().join("truncated.metadata.json"); + fs::write(&metadata_location, fixture).unwrap(); + + let file_io = FileIO::new_with_fs(); + let metadata_location = metadata_location.to_str().unwrap(); + let err = TableMetadata::read_from(&file_io, metadata_location) + .await + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + assert!(err.to_string().contains(metadata_location)); + } + #[tokio::test] async fn test_table_metadata_read_nonexistent_file() { // Create a FileIO instance @@ -3750,6 +3828,47 @@ mod tests { assert_eq!(read_metadata, compressed_metadata); } + #[tokio::test] + async fn test_table_metadata_write_with_zstd_compression() { + let temp_dir = TempDir::new().unwrap(); + let temp_path = temp_dir.path().to_str().unwrap(); + let file_io = FileIO::new_with_fs(); + let original_metadata: TableMetadata = + get_test_table_metadata_at("TableMetadataV2Valid.json", temp_path); + + let mut props = original_metadata.properties.clone(); + props.insert( + TableProperties::PROPERTY_METADATA_COMPRESSION_CODEC.to_string(), + "ZsTd".to_string(), + ); + let compressed_metadata = + TableMetadataBuilder::new_from_metadata(original_metadata.clone(), None) + .assign_uuid(original_metadata.table_uuid) + .set_properties(props) + .unwrap() + .build() + .unwrap() + .metadata; + + let expected_location = + format!("{temp_path}/00000-2cd22b57-5127-4198-92ba-e4e67c79821b.zstd.metadata.json"); + let metadata_location = expected_location.parse::().unwrap(); + + compressed_metadata + .write_to(&file_io, &metadata_location) + .await + .unwrap(); + + assert_eq!(metadata_location.to_string(), expected_location); + let raw_content = fs::read(&expected_location).unwrap(); + assert!(raw_content.starts_with(ZSTD_MAGIC)); + + let read_metadata = TableMetadata::read_from(&file_io, &expected_location) + .await + .unwrap(); + assert_eq!(read_metadata, compressed_metadata); + } + #[test] fn test_partition_name_exists() { let schema = Schema::builder() diff --git a/crates/iceberg/src/spec/table_properties.rs b/crates/iceberg/src/spec/table_properties.rs index 379feee5c1..4c863df741 100644 --- a/crates/iceberg/src/spec/table_properties.rs +++ b/crates/iceberg/src/spec/table_properties.rs @@ -73,7 +73,7 @@ fn parse_location_property( /// Parse compression codec for metadata files from table properties. /// Retrieves the compression codec property, applies defaults, and parses the value. -/// Only "none" (or empty string) and "gzip" are supported for metadata compression. +/// Only "none" (or empty string), "gzip", and "zstd" are supported for metadata compression. /// /// # Arguments /// @@ -81,8 +81,8 @@ fn parse_location_property( /// /// # Errors /// -/// Returns an error if the codec is not "none", "", or "gzip" (case-insensitive). -/// Lz4 and Zstd are not supported for metadata file compression. +/// Returns an error if the codec is not "none", "", "gzip", or "zstd" (case-insensitive). +/// Lz4 and Snappy are not supported for metadata file compression. pub(crate) fn parse_metadata_file_compression( properties: &HashMap, ) -> Result { @@ -107,22 +107,24 @@ pub(crate) fn parse_metadata_file_compression( Error::new( ErrorKind::DataInvalid, format!( - "Invalid metadata compression codec: {value}. Only '{}' and '{}' are supported.", + "Invalid metadata compression codec: {value}. Only '{}', '{}', and '{}' are supported.", CompressionCodec::None.name(), - CompressionCodec::gzip_default().name() + CompressionCodec::gzip_default().name(), + CompressionCodec::zstd_default().name() ), ) })?; - // Validate that only None and Gzip are used for metadata + // Validate that only None, Gzip, and Zstd are used for metadata match codec { - CompressionCodec::None | CompressionCodec::Gzip(_) => Ok(codec), + CompressionCodec::None | CompressionCodec::Gzip(_) | CompressionCodec::Zstd(_) => Ok(codec), _ => Err(Error::new( ErrorKind::DataInvalid, format!( - "Invalid metadata compression codec: {value}. Only '{}' and '{}' are supported for metadata files.", + "Invalid metadata compression codec: {value}. Only '{}', '{}', and '{}' are supported for metadata files.", CompressionCodec::None.name(), - CompressionCodec::gzip_default().name() + CompressionCodec::gzip_default().name(), + CompressionCodec::zstd_default().name() ), )), } @@ -547,15 +549,17 @@ mod tests { #[test] fn test_table_properties_compression() { - let props = HashMap::from([( - TableProperties::PROPERTY_METADATA_COMPRESSION_CODEC.to_string(), - "gzip".to_string(), - )]); - let table_properties = TableProperties::try_from(&props).unwrap(); - assert_eq!( - table_properties.metadata_compression_codec, - CompressionCodec::gzip_default() - ); + for (value, expected) in [ + ("gzip", CompressionCodec::gzip_default()), + ("zstd", CompressionCodec::zstd_default()), + ] { + let props = HashMap::from([( + TableProperties::PROPERTY_METADATA_COMPRESSION_CODEC.to_string(), + value.to_string(), + )]); + let table_properties = TableProperties::try_from(&props).unwrap(); + assert_eq!(table_properties.metadata_compression_codec, expected); + } } #[test] @@ -573,38 +577,20 @@ mod tests { #[test] fn test_table_properties_compression_case_insensitive() { - // Test uppercase - let props_upper = HashMap::from([( - TableProperties::PROPERTY_METADATA_COMPRESSION_CODEC.to_string(), - "GZIP".to_string(), - )]); - let table_properties = TableProperties::try_from(&props_upper).unwrap(); - assert_eq!( - table_properties.metadata_compression_codec, - CompressionCodec::gzip_default() - ); - - // Test mixed case - let props_mixed = HashMap::from([( - TableProperties::PROPERTY_METADATA_COMPRESSION_CODEC.to_string(), - "GzIp".to_string(), - )]); - let table_properties = TableProperties::try_from(&props_mixed).unwrap(); - assert_eq!( - table_properties.metadata_compression_codec, - CompressionCodec::gzip_default() - ); - - // Test "NONE" should also be case-insensitive - let props_none_upper = HashMap::from([( - TableProperties::PROPERTY_METADATA_COMPRESSION_CODEC.to_string(), - "NONE".to_string(), - )]); - let table_properties = TableProperties::try_from(&props_none_upper).unwrap(); - assert_eq!( - table_properties.metadata_compression_codec, - CompressionCodec::None - ); + for (value, expected) in [ + ("GZIP", CompressionCodec::gzip_default()), + ("GzIp", CompressionCodec::gzip_default()), + ("ZSTD", CompressionCodec::zstd_default()), + ("ZsTd", CompressionCodec::zstd_default()), + ("NONE", CompressionCodec::None), + ] { + let props = HashMap::from([( + TableProperties::PROPERTY_METADATA_COMPRESSION_CODEC.to_string(), + value.to_string(), + )]); + let table_properties = TableProperties::try_from(&props).unwrap(); + assert_eq!(table_properties.metadata_compression_codec, expected); + } } #[test] @@ -698,7 +684,7 @@ mod tests { #[test] fn test_table_properties_compression_invalid_rejected() { - let invalid_codecs = ["lz4", "zstd", "snappy"]; + let invalid_codecs = ["lz4", "snappy"]; for codec in invalid_codecs { let props = HashMap::from([( @@ -712,7 +698,7 @@ mod tests { "Expected error message to contain codec '{codec}', got: {err_msg}" ); assert!( - err_msg.contains("Only 'none' and 'gzip' are supported"), + err_msg.contains("Only 'none', 'gzip', and 'zstd' are supported"), "Expected error message to contain supported codecs, got: {err_msg}" ); } @@ -720,65 +706,23 @@ mod tests { #[test] fn test_parse_metadata_file_compression_valid() { - // Test with "none" - let props = HashMap::from([( - TableProperties::PROPERTY_METADATA_COMPRESSION_CODEC.to_string(), - "none".to_string(), - )]); - assert_eq!( - parse_metadata_file_compression(&props).unwrap(), - CompressionCodec::None - ); - - // Test with empty string - let props = HashMap::from([( - TableProperties::PROPERTY_METADATA_COMPRESSION_CODEC.to_string(), - "".to_string(), - )]); - assert_eq!( - parse_metadata_file_compression(&props).unwrap(), - CompressionCodec::None - ); - - // Test with "gzip" - let props = HashMap::from([( - TableProperties::PROPERTY_METADATA_COMPRESSION_CODEC.to_string(), - "gzip".to_string(), - )]); - assert_eq!( - parse_metadata_file_compression(&props).unwrap(), - CompressionCodec::gzip_default() - ); - - // Test case insensitivity - "NONE" - let props = HashMap::from([( - TableProperties::PROPERTY_METADATA_COMPRESSION_CODEC.to_string(), - "NONE".to_string(), - )]); - assert_eq!( - parse_metadata_file_compression(&props).unwrap(), - CompressionCodec::None - ); - - // Test case insensitivity - "GZIP" - let props = HashMap::from([( - TableProperties::PROPERTY_METADATA_COMPRESSION_CODEC.to_string(), - "GZIP".to_string(), - )]); - assert_eq!( - parse_metadata_file_compression(&props).unwrap(), - CompressionCodec::gzip_default() - ); - - // Test case insensitivity - "GzIp" - let props = HashMap::from([( - TableProperties::PROPERTY_METADATA_COMPRESSION_CODEC.to_string(), - "GzIp".to_string(), - )]); - assert_eq!( - parse_metadata_file_compression(&props).unwrap(), - CompressionCodec::gzip_default() - ); + for (value, expected) in [ + ("none", CompressionCodec::None), + ("", CompressionCodec::None), + ("gzip", CompressionCodec::gzip_default()), + ("zstd", CompressionCodec::zstd_default()), + ("NONE", CompressionCodec::None), + ("GZIP", CompressionCodec::gzip_default()), + ("GzIp", CompressionCodec::gzip_default()), + ("ZSTD", CompressionCodec::zstd_default()), + ("ZsTd", CompressionCodec::zstd_default()), + ] { + let props = HashMap::from([( + TableProperties::PROPERTY_METADATA_COMPRESSION_CODEC.to_string(), + value.to_string(), + )]); + assert_eq!(parse_metadata_file_compression(&props).unwrap(), expected); + } // Test default when property is missing let props = HashMap::new(); @@ -790,7 +734,7 @@ mod tests { #[test] fn test_parse_metadata_file_compression_invalid() { - let invalid_codecs = ["lz4", "zstd", "snappy"]; + let invalid_codecs = ["lz4", "snappy"]; for codec in invalid_codecs { let props = HashMap::from([( @@ -804,7 +748,7 @@ mod tests { "Expected error message to contain 'Invalid metadata compression codec', got: {err_msg}" ); assert!( - err_msg.contains("Only 'none' and 'gzip' are supported"), + err_msg.contains("Only 'none', 'gzip', and 'zstd' are supported"), "Expected error message to contain supported codecs, got: {err_msg}" ); } diff --git a/crates/iceberg/testdata/table_metadata/TableMetadataV2Valid.zstd.metadata.json b/crates/iceberg/testdata/table_metadata/TableMetadataV2Valid.zstd.metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..4709076c12c14f984e4d643122cc614a4139e498 GIT binary patch literal 374 zcmV-+0g3)7wJ-fd0K)-23jp#KK_FmBs{m*i3bii31=Ja3F^QG18?w zi+k5HBM-Rw4y(uWfF(i*vSuB3Ii8yhs?&1Y^^pD`laUcluZvJoyhL6Ro{mtSZgUl0 zSK~A+kjp+T(ic1HxIoF%k7JzbOeFwsEB1s7Y=|)htHTE=kg^EW5YV9oT#(`w%VE3( zfQC50e225MljHNaNxui3!Ru^`N5eky>*M?B36JyCaNmvcm$DNL%l Un{Rq*z$$7xubJCy0yLRcBcl4NX8-^I literal 0 HcmV?d00001 From fdc645c14ad51347fdd1ec3b8abf6760a4c03360 Mon Sep 17 00:00:00 2001 From: xuanyili Date: Tue, 28 Jul 2026 13:46:50 +0000 Subject: [PATCH 2/2] refactor: centralize metadata compression mappings Reuse shared supported-codec, suffix, and magic mappings across metadata property validation, location parsing, and I/O. Simplify alternate suffix parsing with a dedicated helper. --- .../iceberg/src/catalog/metadata_location.rs | 57 ++++++++--------- crates/iceberg/src/compression.rs | 24 ++++++- crates/iceberg/src/spec/table_metadata.rs | 47 ++++++-------- crates/iceberg/src/spec/table_properties.rs | 63 ++++++++++--------- 4 files changed, 104 insertions(+), 87 deletions(-) diff --git a/crates/iceberg/src/catalog/metadata_location.rs b/crates/iceberg/src/catalog/metadata_location.rs index 025fbda426..d8891d349e 100644 --- a/crates/iceberg/src/catalog/metadata_location.rs +++ b/crates/iceberg/src/catalog/metadata_location.rs @@ -21,13 +21,25 @@ use std::str::FromStr; use uuid::Uuid; -use crate::compression::CompressionCodec; +use crate::compression::{CompressionCodec, TABLE_METADATA_SUFFIX_TO_COMPRESSION}; use crate::spec::{TableMetadata, parse_metadata_file_compression}; use crate::{Error, ErrorKind, Result}; /// Default folder name for metadata files under the table location, used when the /// `write.metadata.path` table property is not set. pub(crate) const METADATA_FOLDER_NAME: &str = "metadata"; +const METADATA_SUFFIX: &str = ".metadata.json"; + +fn strip_metadata_suffix<'a>(file_name: &'a str, codec_suffix: &str) -> Option<&'a str> { + file_name + .strip_suffix(METADATA_SUFFIX) + .and_then(|stripped| stripped.strip_suffix(codec_suffix)) + .or_else(|| { + file_name + .strip_suffix(codec_suffix) + .and_then(|stripped| stripped.strip_suffix(METADATA_SUFFIX)) + }) +} /// Helper for parsing a location of the format: `/-.metadata.json` /// or with compression: `/-..metadata.json` @@ -91,34 +103,23 @@ impl MetadataLocation { /// Parses a file name of the format `-.metadata.json` /// or with compression before or after `.metadata.json`. - /// Parse errors for compression codec result in CompressionCodec::None. fn parse_file_name(file_name: &str) -> Result<(i32, Uuid, CompressionCodec)> { - let gzip_suffix = CompressionCodec::gzip_default().suffix()?; - let zstd_suffix = CompressionCodec::zstd_default().suffix()?; - let metadata_suffix = ".metadata.json"; - - let (stripped, compression_codec) = if let Some(stripped) = - file_name.strip_suffix(&format!("{metadata_suffix}{zstd_suffix}")) - { - (stripped, CompressionCodec::zstd_default()) - } else if let Some(stripped) = - file_name.strip_suffix(&format!("{metadata_suffix}{gzip_suffix}")) - { - (stripped, CompressionCodec::gzip_default()) - } else if let Some(stripped) = file_name.strip_suffix(metadata_suffix) { - if let Some(stripped) = stripped.strip_suffix(zstd_suffix) { - (stripped, CompressionCodec::zstd_default()) - } else if let Some(stripped) = stripped.strip_suffix(gzip_suffix) { - (stripped, CompressionCodec::gzip_default()) - } else { - (stripped, CompressionCodec::None) - } - } else { - return Err(Error::new( - ErrorKind::Unexpected, - format!("Invalid metadata file ending: {file_name}"), - )); - }; + let (stripped, compression_codec) = TABLE_METADATA_SUFFIX_TO_COMPRESSION + .iter() + .find_map(|(suffix, codec)| { + strip_metadata_suffix(file_name, suffix).map(|stripped| (stripped, *codec)) + }) + .or_else(|| { + file_name + .strip_suffix(METADATA_SUFFIX) + .map(|stripped| (stripped, CompressionCodec::None)) + }) + .ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + format!("Invalid metadata file ending: {file_name}"), + ) + })?; let (version, id) = stripped.split_once('-').ok_or(Error::new( ErrorKind::Unexpected, diff --git a/crates/iceberg/src/compression.rs b/crates/iceberg/src/compression.rs index 266d93fbe8..fdd0524de4 100644 --- a/crates/iceberg/src/compression.rs +++ b/crates/iceberg/src/compression.rs @@ -33,6 +33,10 @@ const ZSTD_DEFAULT_LEVEL: u8 = 3; const GZIP_DEFAULT_LEVEL: u8 = 6; /// Maximum compression level for Gzip. const GZIP_MAX_LEVEL: u8 = 9; +const GZIP_SUFFIX: &str = ".gz"; +const ZSTD_SUFFIX: &str = ".zstd"; +const GZIP_MAGIC: &[u8] = &[0x1F, 0x8B]; +const ZSTD_MAGIC: &[u8] = &[0x28, 0xB5, 0x2F, 0xFD]; /// Data compression formats #[derive(Debug, PartialEq, Eq, Clone, Copy, Default)] @@ -53,6 +57,22 @@ pub enum CompressionCodec { Snappy, } +pub(crate) const TABLE_METADATA_SUPPORTED_COMPRESSION: &[CompressionCodec] = &[ + CompressionCodec::None, + CompressionCodec::Gzip(GZIP_DEFAULT_LEVEL), + CompressionCodec::Zstd(ZSTD_DEFAULT_LEVEL), +]; + +pub(crate) const TABLE_METADATA_SUFFIX_TO_COMPRESSION: &[(&str, CompressionCodec)] = &[ + (ZSTD_SUFFIX, CompressionCodec::Zstd(ZSTD_DEFAULT_LEVEL)), + (GZIP_SUFFIX, CompressionCodec::Gzip(GZIP_DEFAULT_LEVEL)), +]; + +pub(crate) const TABLE_METADATA_MAGIC_TO_COMPRESSION: &[(&[u8], CompressionCodec)] = &[ + (GZIP_MAGIC, CompressionCodec::Gzip(GZIP_DEFAULT_LEVEL)), + (ZSTD_MAGIC, CompressionCodec::Zstd(ZSTD_DEFAULT_LEVEL)), +]; + impl CompressionCodec { /// Returns a Zstd codec with the default compression level. pub const fn zstd_default() -> Self { @@ -177,8 +197,8 @@ impl CompressionCodec { pub fn suffix(&self) -> Result<&'static str> { match self { CompressionCodec::None => Ok(""), - CompressionCodec::Zstd(_) => Ok(".zstd"), - CompressionCodec::Gzip(_) => Ok(".gz"), + CompressionCodec::Zstd(_) => Ok(ZSTD_SUFFIX), + CompressionCodec::Gzip(_) => Ok(GZIP_SUFFIX), codec @ (CompressionCodec::Lz4 | CompressionCodec::Snappy) => Err(Error::new( ErrorKind::FeatureUnsupported, format!("suffix not defined for {codec:?}"), diff --git a/crates/iceberg/src/spec/table_metadata.rs b/crates/iceberg/src/spec/table_metadata.rs index 0733bac1f6..e89b580e5b 100644 --- a/crates/iceberg/src/spec/table_metadata.rs +++ b/crates/iceberg/src/spec/table_metadata.rs @@ -38,7 +38,7 @@ use super::{ TableProperties, parse_metadata_file_compression, }; use crate::catalog::{METADATA_FOLDER_NAME, MetadataLocation}; -use crate::compression::CompressionCodec; +use crate::compression::{CompressionCodec, TABLE_METADATA_MAGIC_TO_COMPRESSION}; use crate::error::{Result, timestamp_ms_to_utc}; use crate::io::FileIO; use crate::spec::EncryptedKey; @@ -46,8 +46,6 @@ use crate::{Error, ErrorKind}; static MAIN_BRANCH: &str = "main"; pub(crate) static ONE_MINUTE_MS: i64 = 60_000; -const GZIP_MAGIC: &[u8] = &[0x1F, 0x8B]; -const ZSTD_MAGIC: &[u8] = &[0x28, 0xB5, 0x2F, 0xFD]; /// Sentinel value used by the Java implementation and older metadata files /// to represent a missing/empty current snapshot ID. During deserialization, @@ -468,13 +466,9 @@ impl TableMetadata { let input_file = file_io.new_input(metadata_location)?; let metadata_content = input_file.read().await?; - let compression_codec = if metadata_content.starts_with(GZIP_MAGIC) { - Some(CompressionCodec::gzip_default()) - } else if metadata_content.starts_with(ZSTD_MAGIC) { - Some(CompressionCodec::zstd_default()) - } else { - None - }; + let compression_codec = TABLE_METADATA_MAGIC_TO_COMPRESSION + .iter() + .find_map(|(magic, codec)| metadata_content.starts_with(magic).then_some(*codec)); let metadata = if let Some(codec) = compression_codec { let decompressed_data = codec.decompress(metadata_content.to_vec()).map_err(|e| { @@ -515,17 +509,7 @@ impl TableMetadata { )); } - // Apply compression based on codec - let data_to_write = match codec { - CompressionCodec::Gzip(_) | CompressionCodec::Zstd(_) => codec.compress(json_data)?, - CompressionCodec::None => json_data, - _ => { - return Err(Error::new( - ErrorKind::DataInvalid, - format!("Unsupported metadata compression codec: {codec:?}"), - )); - } - }; + let data_to_write = codec.compress(json_data)?; file_io .new_output(metadata_location.to_string())? @@ -1633,9 +1617,9 @@ mod tests { use tempfile::TempDir; use uuid::Uuid; - use super::{FormatVersion, MetadataLog, SnapshotLog, TableMetadataBuilder, ZSTD_MAGIC}; + use super::{FormatVersion, MetadataLog, SnapshotLog, TableMetadataBuilder}; use crate::catalog::MetadataLocation; - use crate::compression::CompressionCodec; + use crate::compression::{CompressionCodec, TABLE_METADATA_MAGIC_TO_COMPRESSION}; use crate::io::FileIO; use crate::spec::table_metadata::TableMetadata; use crate::spec::{ @@ -1663,6 +1647,13 @@ mod tests { serde_json::from_str(&metadata).unwrap() } + fn table_metadata_magic(codec: CompressionCodec) -> &'static [u8] { + TABLE_METADATA_MAGIC_TO_COMPRESSION + .iter() + .find_map(|(magic, mapped_codec)| (*mapped_codec == codec).then_some(*magic)) + .unwrap() + } + /// Loads a test table metadata and relocates it to `location`, so that derived /// metadata paths point at a writable (e.g. temp) directory. fn get_test_table_metadata_at(file_name: &str, location: &str) -> TableMetadata { @@ -3688,7 +3679,7 @@ mod tests { async fn test_table_metadata_read_iceberg_go_zstd_fixture() { let fixture = fs::read("testdata/table_metadata/TableMetadataV2Valid.zstd.metadata.json").unwrap(); - assert!(fixture.starts_with(ZSTD_MAGIC)); + assert!(fixture.starts_with(table_metadata_magic(CompressionCodec::zstd_default()))); let temp_dir = TempDir::new().unwrap(); let file_io = FileIO::new_with_fs(); @@ -3743,7 +3734,7 @@ mod tests { let mut fixture = fs::read("testdata/table_metadata/TableMetadataV2Valid.zstd.metadata.json").unwrap(); fixture.truncate(fixture.len() / 2); - assert!(fixture.starts_with(ZSTD_MAGIC)); + assert!(fixture.starts_with(table_metadata_magic(CompressionCodec::zstd_default()))); let temp_dir = TempDir::new().unwrap(); let metadata_location = temp_dir.path().join("truncated.metadata.json"); @@ -3815,9 +3806,7 @@ mod tests { // Read the raw file and check it's gzip compressed let raw_content = fs::read(&metadata_location_str).unwrap(); - assert!(raw_content.len() > 2); - assert_eq!(raw_content[0], 0x1F); // gzip magic number - assert_eq!(raw_content[1], 0x8B); // gzip magic number + assert!(raw_content.starts_with(table_metadata_magic(CompressionCodec::gzip_default()))); // Read the metadata back using the compressed location let read_metadata = TableMetadata::read_from(&file_io, &metadata_location_str) @@ -3861,7 +3850,7 @@ mod tests { assert_eq!(metadata_location.to_string(), expected_location); let raw_content = fs::read(&expected_location).unwrap(); - assert!(raw_content.starts_with(ZSTD_MAGIC)); + assert!(raw_content.starts_with(table_metadata_magic(CompressionCodec::zstd_default()))); let read_metadata = TableMetadata::read_from(&file_io, &expected_location) .await diff --git a/crates/iceberg/src/spec/table_properties.rs b/crates/iceberg/src/spec/table_properties.rs index 4c863df741..ee71c27dae 100644 --- a/crates/iceberg/src/spec/table_properties.rs +++ b/crates/iceberg/src/spec/table_properties.rs @@ -19,9 +19,35 @@ use std::collections::HashMap; use std::fmt::Display; use std::str::FromStr; -use crate::compression::CompressionCodec; +use crate::compression::{CompressionCodec, TABLE_METADATA_SUPPORTED_COMPRESSION}; use crate::error::{Error, ErrorKind, Result}; +fn supported_metadata_compression_names() -> String { + let names = TABLE_METADATA_SUPPORTED_COMPRESSION + .iter() + .map(|codec| format!("'{}'", codec.name())) + .collect::>(); + let (last, rest) = names + .split_last() + .expect("metadata compression codec list must not be empty"); + + if rest.is_empty() { + last.clone() + } else { + format!("{}, and {last}", rest.join(", ")) + } +} + +fn invalid_metadata_compression_codec(value: &str) -> Error { + Error::new( + ErrorKind::DataInvalid, + format!( + "Invalid metadata compression codec: {value}. Only {} are supported for metadata files.", + supported_metadata_compression_names() + ), + ) +} + fn parse_property( properties: &HashMap, key: &str, @@ -100,33 +126,14 @@ pub(crate) fn parse_metadata_file_compression( let lowercase_value = value.to_lowercase(); // Use serde to parse the codec (which has rename_all = "lowercase") - let codec: CompressionCodec = serde_json::from_value(serde_json::Value::String( - lowercase_value, - )) - .map_err(|_| { - Error::new( - ErrorKind::DataInvalid, - format!( - "Invalid metadata compression codec: {value}. Only '{}', '{}', and '{}' are supported.", - CompressionCodec::None.name(), - CompressionCodec::gzip_default().name(), - CompressionCodec::zstd_default().name() - ), - ) - })?; - - // Validate that only None, Gzip, and Zstd are used for metadata - match codec { - CompressionCodec::None | CompressionCodec::Gzip(_) | CompressionCodec::Zstd(_) => Ok(codec), - _ => Err(Error::new( - ErrorKind::DataInvalid, - format!( - "Invalid metadata compression codec: {value}. Only '{}', '{}', and '{}' are supported for metadata files.", - CompressionCodec::None.name(), - CompressionCodec::gzip_default().name(), - CompressionCodec::zstd_default().name() - ), - )), + let codec: CompressionCodec = + serde_json::from_value(serde_json::Value::String(lowercase_value)) + .map_err(|_| invalid_metadata_compression_codec(value))?; + + if TABLE_METADATA_SUPPORTED_COMPRESSION.contains(&codec) { + Ok(codec) + } else { + Err(invalid_metadata_compression_codec(value)) } }