Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
202 changes: 149 additions & 53 deletions crates/iceberg/src/catalog/metadata_location.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,16 +21,28 @@ 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-dir>/<version>-<uuid>.metadata.json`
/// or with compression: `<metadata-dir>/<version>-<uuid>.gz.metadata.json`
/// or with compression: `<metadata-dir>/<version>-<uuid>.<codec>.metadata.json`
///
/// `<metadata-dir>` is set to the `write.metadata.path` table property and
/// it defaults to the `<location>/metadata` when the property is not set.
Expand Down Expand Up @@ -90,21 +102,24 @@ impl MetadataLocation {
}

/// Parses a file name of the format `<version>-<uuid>.metadata.json`
/// or with compression: `<version>-<uuid>.gz.metadata.json`.
/// Parse errors for compression codec result in CompressionCodec::None.
/// or with compression before or after `.metadata.json`.
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())
} else {
(stripped, CompressionCodec::None)
};
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,
Expand Down Expand Up @@ -248,6 +263,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",
Expand Down Expand Up @@ -295,6 +340,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());
Expand All @@ -318,33 +382,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]
Expand Down Expand Up @@ -374,10 +439,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!(
Expand Down Expand Up @@ -436,15 +517,30 @@ mod test {

#[test]
fn test_new_with_metadata_honors_write_metadata_path() {
// Test metadata lives under `<location>/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 `<location>/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([(
Expand Down
36 changes: 26 additions & 10 deletions crates/iceberg/src/compression.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand All @@ -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 {
Expand Down Expand Up @@ -169,18 +189,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::Gzip(_) => Ok(".gz"),
codec @ (CompressionCodec::Lz4
| CompressionCodec::Zstd(_)
| CompressionCodec::Snappy) => Err(Error::new(
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:?}"),
)),
Expand Down Expand Up @@ -244,17 +263,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]
Expand Down
Loading
Loading