diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index bb4d350e5a..3cd662add1 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -1086,6 +1086,7 @@ pub const iceberg::io::S3_ASSUME_ROLE_SESSION_NAME: &str pub const iceberg::io::S3_DISABLE_CONFIG_LOAD: &str pub const iceberg::io::S3_DISABLE_EC2_METADATA: &str pub const iceberg::io::S3_ENDPOINT: &str +pub const iceberg::io::S3_MULTIPART_PART_SIZE_BYTES: &str pub const iceberg::io::S3_PATH_STYLE_ACCESS: &str pub const iceberg::io::S3_REGION: &str pub const iceberg::io::S3_SECRET_ACCESS_KEY: &str diff --git a/crates/iceberg/src/io/storage/config/s3.rs b/crates/iceberg/src/io/storage/config/s3.rs index ae3bc4ba3b..82f27e15dd 100644 --- a/crates/iceberg/src/io/storage/config/s3.rs +++ b/crates/iceberg/src/io/storage/config/s3.rs @@ -65,6 +65,16 @@ pub const S3_ALLOW_ANONYMOUS: &str = "s3.allow-anonymous"; pub const S3_DISABLE_EC2_METADATA: &str = "s3.disable-ec2-metadata"; /// Option to skip loading configuration from config file and the env. pub const S3_DISABLE_CONFIG_LOAD: &str = "s3.disable-config-load"; +/// Size in bytes of each part of a multipart upload. Defaults to 32 MiB, matching Java +/// `S3FileIOProperties.MULTIPART_SIZE_DEFAULT`. Must be at least 5 MiB, the S3 minimum, +/// and at most `i32::MAX`, since Java reads the property as an `int`. +/// +/// Each part travels as one request, so the part size bounds how much data has to +/// transfer within the per-operation IO timeout. Lower it on a slow connection; if a +/// 5 MiB part still cannot finish in time, raise `opendal.io-timeout-ms` instead. +/// +/// Only `iceberg-storage-opendal` reads this property, so [`S3Config`] has no matching field. +pub const S3_MULTIPART_PART_SIZE_BYTES: &str = "s3.multipart.part-size-bytes"; /// Amazon S3 storage configuration. /// diff --git a/crates/storage/opendal/public-api.txt b/crates/storage/opendal/public-api.txt index c2e310335c..f5e3baa817 100644 --- a/crates/storage/opendal/public-api.txt +++ b/crates/storage/opendal/public-api.txt @@ -19,10 +19,11 @@ pub iceberg_storage_opendal::OpenDalStorage::Memory::operator: opendal_core::typ pub iceberg_storage_opendal::OpenDalStorage::Oss pub iceberg_storage_opendal::OpenDalStorage::Oss::client_config: iceberg_storage_opendal::OpenDalClientConfig pub iceberg_storage_opendal::OpenDalStorage::Oss::config: alloc::sync::Arc -pub iceberg_storage_opendal::OpenDalStorage::S3 +#[non_exhaustive] pub iceberg_storage_opendal::OpenDalStorage::S3 pub iceberg_storage_opendal::OpenDalStorage::S3::client_config: iceberg_storage_opendal::OpenDalClientConfig pub iceberg_storage_opendal::OpenDalStorage::S3::config: alloc::sync::Arc pub iceberg_storage_opendal::OpenDalStorage::S3::customized_credential_load: core::option::Option +pub iceberg_storage_opendal::OpenDalStorage::S3::multipart_part_size: u64 impl core::clone::Clone for iceberg_storage_opendal::OpenDalStorage pub fn iceberg_storage_opendal::OpenDalStorage::clone(&self) -> iceberg_storage_opendal::OpenDalStorage impl core::fmt::Debug for iceberg_storage_opendal::OpenDalStorage diff --git a/crates/storage/opendal/src/lib.rs b/crates/storage/opendal/src/lib.rs index 9e01552016..a13d380fae 100644 --- a/crates/storage/opendal/src/lib.rs +++ b/crates/storage/opendal/src/lib.rs @@ -43,6 +43,7 @@ use iceberg::{Error, ErrorKind, Result}; use iceberg_property_macro::Properties; use opendal::Operator; use opendal::layers::{RetryLayer, TimeoutLayer}; +use opendal::options::WriteOptions; use serde::{Deserialize, Serialize}; use utils::from_opendal_error; @@ -228,6 +229,7 @@ impl StorageFactory for OpenDalStorageFactory { customized_credential_load, } => Ok(Arc::new(OpenDalStorage::S3 { config: s3_config_parse(config.props().clone())?.into(), + multipart_part_size: s3_multipart_part_size_parse(config.props())?, customized_credential_load: customized_credential_load.clone(), client_config, })), @@ -301,9 +303,13 @@ pub enum OpenDalStorage { /// Accepts any S3-family URL (`s3://`, `s3a://`, `s3n://`); the scheme is /// derived from the path at call time. #[cfg(feature = "opendal-s3")] + #[non_exhaustive] S3 { /// S3 configuration. config: Arc, + /// Bytes carried by one multipart upload request. + #[serde(default = "default_multipart_part_size")] + multipart_part_size: u64, /// Custom AWS credential loader. #[serde(skip)] customized_credential_load: Option, @@ -513,6 +519,24 @@ impl OpenDalStorage { } } + /// Bounds the S3 request size. Without a chunk size OpenDAL turns each caller + /// buffer into one request, and `ParquetWriter` hands over a whole row group + /// at a time. Other backends keep OpenDAL's defaults. + #[allow(unreachable_patterns)] + fn write_options(&self) -> WriteOptions { + match self { + #[cfg(feature = "opendal-s3")] + OpenDalStorage::S3 { + multipart_part_size, + .. + } => WriteOptions { + chunk: Some(usize::try_from(*multipart_part_size).unwrap_or(usize::MAX)), + ..WriteOptions::default() + }, + _ => WriteOptions::default(), + } + } + /// Returns a cache key used by `delete_stream` to group paths by storage operator. /// /// For most backends the URL host (bucket name) is sufficient. For HF the host @@ -662,7 +686,7 @@ impl Storage for OpenDalStorage { async fn write(&self, path: &str, bs: Bytes) -> Result<()> { let (op, relative_path) = self.create_operator(&path)?; - op.write(relative_path, bs) + op.write_options(relative_path, bs, self.write_options()) .await .map_err(from_opendal_error)?; Ok(()) @@ -671,7 +695,9 @@ impl Storage for OpenDalStorage { async fn writer(&self, path: &str) -> Result> { let (op, relative_path) = self.create_operator(&path)?; Ok(Box::new(OpenDalWriter::new( - op.writer(relative_path).await.map_err(from_opendal_error)?, + op.writer_options(relative_path, self.write_options()) + .await + .map_err(from_opendal_error)?, ))) } @@ -866,6 +892,7 @@ mod tests { fn test_client_config_serde_round_trip() { let storage = OpenDalStorage::S3 { config: Arc::new(S3Config::default()), + multipart_part_size: default_multipart_part_size(), customized_credential_load: None, client_config: client_config("45000").unwrap(), }; @@ -1080,6 +1107,7 @@ mod tests { fn test_relativize_path_s3() { let storage = OpenDalStorage::S3 { config: Arc::new(S3Config::default()), + multipart_part_size: default_multipart_part_size(), customized_credential_load: None, client_config: OpenDalClientConfig::default(), }; diff --git a/crates/storage/opendal/src/resolving.rs b/crates/storage/opendal/src/resolving.rs index 51e76da472..73b39eea2c 100644 --- a/crates/storage/opendal/src/resolving.rs +++ b/crates/storage/opendal/src/resolving.rs @@ -93,6 +93,7 @@ fn build_storage_for_scheme( let config = crate::s3::s3_config_parse(props.clone())?; Ok(OpenDalStorage::S3 { config: Arc::new(config), + multipart_part_size: crate::s3::s3_multipart_part_size_parse(props)?, customized_credential_load: customized_credential_load.clone(), client_config, }) diff --git a/crates/storage/opendal/src/s3.rs b/crates/storage/opendal/src/s3.rs index 2997fa81d9..8f1828dc46 100644 --- a/crates/storage/opendal/src/s3.rs +++ b/crates/storage/opendal/src/s3.rs @@ -21,8 +21,8 @@ use std::sync::Arc; use iceberg::io::{ CLIENT_REGION, S3_ACCESS_KEY_ID, S3_ALLOW_ANONYMOUS, S3_ASSUME_ROLE_ARN, S3_ASSUME_ROLE_EXTERNAL_ID, S3_ASSUME_ROLE_SESSION_NAME, S3_DISABLE_CONFIG_LOAD, - S3_DISABLE_EC2_METADATA, S3_ENDPOINT, S3_PATH_STYLE_ACCESS, S3_REGION, S3_SECRET_ACCESS_KEY, - S3_SESSION_TOKEN, S3_SSE_KEY, S3_SSE_MD5, S3_SSE_TYPE, + S3_DISABLE_EC2_METADATA, S3_ENDPOINT, S3_MULTIPART_PART_SIZE_BYTES, S3_PATH_STYLE_ACCESS, + S3_REGION, S3_SECRET_ACCESS_KEY, S3_SESSION_TOKEN, S3_SSE_KEY, S3_SSE_MD5, S3_SSE_TYPE, }; use iceberg::{Error, ErrorKind, Result}; use opendal::services::S3Config; @@ -36,6 +36,50 @@ use url::Url; use crate::utils::{from_opendal_error, is_truthy}; +/// S3 rejects a non-final part smaller than this. +const MULTIPART_PART_SIZE_MIN: u64 = 5 * 1024 * 1024; + +/// Java reads this property as an `int`, so a larger value fails there. +const MULTIPART_PART_SIZE_MAX: u64 = i32::MAX as u64; + +/// Matches Java `S3FileIOProperties.MULTIPART_SIZE_DEFAULT`. +const MULTIPART_PART_SIZE_DEFAULT: u64 = 32 * 1024 * 1024; + +pub(crate) fn default_multipart_part_size() -> u64 { + MULTIPART_PART_SIZE_DEFAULT +} + +/// Parse iceberg props to s3 multipart upload part size. +pub(crate) fn s3_multipart_part_size_parse(m: &HashMap) -> Result { + let Some(value) = m.get(S3_MULTIPART_PART_SIZE_BYTES) else { + return Ok(default_multipart_part_size()); + }; + let part_size = value.parse::().map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + format!("Invalid {S3_MULTIPART_PART_SIZE_BYTES}: {value}: {e}"), + ) + })?; + if part_size < MULTIPART_PART_SIZE_MIN { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "Invalid {S3_MULTIPART_PART_SIZE_BYTES}: {part_size} bytes is below the S3 minimum part size of {} MiB", + MULTIPART_PART_SIZE_MIN / (1024 * 1024) + ), + )); + } + if part_size > MULTIPART_PART_SIZE_MAX { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "Invalid {S3_MULTIPART_PART_SIZE_BYTES}: {part_size} bytes is above {MULTIPART_PART_SIZE_MAX} bytes, the largest value Java accepts" + ), + )); + } + Ok(part_size) +} + /// Parse iceberg props to s3 config. pub(crate) fn s3_config_parse(mut m: HashMap) -> Result { let mut cfg = S3Config::default(); @@ -184,9 +228,12 @@ impl CustomAwsCredentialLoader { mod tests { use std::collections::HashMap; - use iceberg::io::S3_PATH_STYLE_ACCESS; + use iceberg::io::{S3_MULTIPART_PART_SIZE_BYTES, S3_PATH_STYLE_ACCESS}; - use super::s3_config_parse; + use super::{ + MULTIPART_PART_SIZE_DEFAULT, MULTIPART_PART_SIZE_MAX, MULTIPART_PART_SIZE_MIN, Result, + s3_config_parse, s3_multipart_part_size_parse, + }; fn parse_with(prop: Option<&str>) -> bool { let mut props = HashMap::new(); @@ -203,4 +250,31 @@ mod tests { assert!(parse_with(Some("false"))); assert!(!parse_with(Some("true"))); } + + fn parse_part_size(prop: Option<&str>) -> Result { + let mut props = HashMap::new(); + if let Some(v) = prop { + props.insert(S3_MULTIPART_PART_SIZE_BYTES.to_string(), v.to_string()); + } + s3_multipart_part_size_parse(&props) + } + + #[test] + fn s3_multipart_part_size_parse_defaults_and_overrides() { + // Match Iceberg S3FileIOProperties.MULTIPART_SIZE_DEFAULT = 32 MiB. + assert_eq!(MULTIPART_PART_SIZE_DEFAULT, 32 * 1024 * 1024); + assert_eq!(parse_part_size(None).unwrap(), MULTIPART_PART_SIZE_DEFAULT); + assert_eq!(parse_part_size(Some("8388608")).unwrap(), 8 * 1024 * 1024); + } + + #[test] + fn s3_multipart_part_size_parse_rejects_invalid() { + // Match Iceberg S3FileIOProperties.MULTIPART_SIZE_MIN = 5 MiB. + assert_eq!(MULTIPART_PART_SIZE_MIN, 5 * 1024 * 1024); + assert!(parse_part_size(Some(&MULTIPART_PART_SIZE_MIN.to_string())).is_ok()); + assert!(parse_part_size(Some(&(MULTIPART_PART_SIZE_MIN - 1).to_string())).is_err()); + assert!(parse_part_size(Some(&MULTIPART_PART_SIZE_MAX.to_string())).is_ok()); + assert!(parse_part_size(Some(&(MULTIPART_PART_SIZE_MAX + 1).to_string())).is_err()); + assert!(parse_part_size(Some("not-a-number")).is_err()); + } } diff --git a/crates/storage/opendal/tests/file_io_s3_test.rs b/crates/storage/opendal/tests/file_io_s3_test.rs index ac6342dddb..14e54dcf59 100644 --- a/crates/storage/opendal/tests/file_io_s3_test.rs +++ b/crates/storage/opendal/tests/file_io_s3_test.rs @@ -26,33 +26,80 @@ mod tests { use bytes::Bytes; use futures::StreamExt; use iceberg::io::{ - FileIO, FileIOBuilder, S3_ACCESS_KEY_ID, S3_ENDPOINT, S3_PATH_STYLE_ACCESS, S3_REGION, - S3_SECRET_ACCESS_KEY, + FileIO, FileIOBuilder, S3_ACCESS_KEY_ID, S3_ENDPOINT, S3_MULTIPART_PART_SIZE_BYTES, + S3_PATH_STYLE_ACCESS, S3_REGION, S3_SECRET_ACCESS_KEY, }; use iceberg_storage_opendal::{ AwsCredential, CustomAwsCredentialLoader, OpenDalStorageFactory, ProvideCredential, }; use iceberg_test_utils::{get_object_store_endpoint, normalize_test_name_with_parts, set_up}; + use opendal::Configurator; use reqsign_core::Context; + const OBJECT_STORE_ACCESS_KEY_ID: &str = "admin"; + const OBJECT_STORE_SECRET_ACCESS_KEY: &str = "password"; + const OBJECT_STORE_REGION: &str = "us-east-1"; + const OBJECT_STORE_BUCKET: &str = "bucket1"; + async fn get_file_io() -> FileIO { + get_file_io_with_props(vec![]).await + } + + async fn get_file_io_with_props(extra: Vec<(&'static str, String)>) -> FileIO { set_up(); let object_store_endpoint = get_object_store_endpoint(); + let mut props = vec![ + (S3_ENDPOINT, object_store_endpoint), + (S3_ACCESS_KEY_ID, OBJECT_STORE_ACCESS_KEY_ID.to_string()), + ( + S3_SECRET_ACCESS_KEY, + OBJECT_STORE_SECRET_ACCESS_KEY.to_string(), + ), + (S3_REGION, OBJECT_STORE_REGION.to_string()), + (S3_PATH_STYLE_ACCESS, "true".to_string()), + ]; + props.extend(extra); + FileIOBuilder::new(Arc::new(OpenDalStorageFactory::S3 { customized_credential_load: None, })) - .with_props(vec![ - (S3_ENDPOINT, object_store_endpoint), - (S3_ACCESS_KEY_ID, "admin".to_string()), - (S3_SECRET_ACCESS_KEY, "password".to_string()), - (S3_REGION, "us-east-1".to_string()), - (S3_PATH_STYLE_ACCESS, "true".to_string()), - ]) + .with_props(props) .build() } + /// Number of upload parts S3 recorded for `key`. A multipart upload reports + /// a `-` ETag; a single-request upload reports a bare MD5. + /// + /// A server-side encrypted object can report an ETag without the suffix, so + /// this helper holds only for the object store fixture these tests run against. + async fn upload_part_count(key: &str) -> usize { + let mut config = opendal::services::S3Config::default(); + config.endpoint = Some(get_object_store_endpoint()); + config.access_key_id = Some(OBJECT_STORE_ACCESS_KEY_ID.to_string()); + config.secret_access_key = Some(OBJECT_STORE_SECRET_ACCESS_KEY.to_string()); + config.region = Some(OBJECT_STORE_REGION.to_string()); + config.bucket = OBJECT_STORE_BUCKET.to_string(); + // Each `#[tokio::test]` has its own runtime, and the first request after + // an earlier one went away fails before the client reconnects. The + // storage under test carries the same layer. + let op = opendal::Operator::new(config.into_builder()) + .unwrap() + .layer(opendal::layers::RetryLayer::new()); + let meta = op.stat(key).await.unwrap(); + let etag = meta + .etag() + .expect("object store reports an ETag") + .trim_matches('"'); + match etag.rsplit_once('-') { + Some((_, parts)) => parts + .parse() + .unwrap_or_else(|e| panic!("unexpected ETag format {etag}: {e}")), + None => 1, + } + } + fn roundtrip_file_io(file_io: &FileIO) -> FileIO { let serialized = file_io.serialize_all().unwrap(); FileIO::deserialize_all(&serialized).unwrap() @@ -62,7 +109,7 @@ mod tests { async fn test_file_io_s3_serialization_roundtrip() { let file_io = roundtrip_file_io(&get_file_io().await); let path = format!( - "s3://bucket1/{}", + "s3://{OBJECT_STORE_BUCKET}/{}", normalize_test_name_with_parts!("test_file_io_s3_serialization_roundtrip") ); @@ -85,7 +132,12 @@ mod tests { async fn test_file_io_s3_exists() { let file_io = get_file_io().await; assert!(!file_io.exists("s3://bucket2/any").await.unwrap()); - assert!(file_io.exists("s3://bucket1/").await.unwrap()); + assert!( + file_io + .exists(&format!("s3://{OBJECT_STORE_BUCKET}/")) + .await + .unwrap() + ); } #[tokio::test] @@ -93,7 +145,7 @@ mod tests { let file_io = get_file_io().await; // Use unique file path based on module path to avoid conflicts let output_path = format!( - "s3://bucket1/{}", + "s3://{OBJECT_STORE_BUCKET}/{}", normalize_test_name_with_parts!("test_file_io_s3_output") ); // Clean up from any previous test runs @@ -106,12 +158,88 @@ mod tests { assert!(file_io.exists(&output_path).await.unwrap()); } + /// `ParquetWriter` hands over a whole row group in one `write` call, so an + /// unbounded part size would make the request as large as that row group. + #[tokio::test] + async fn test_file_io_s3_writer_splits_write_into_bounded_parts() { + const PART_SIZE: usize = 5 * 1024 * 1024; + const PARTS: usize = 3; + + let file_io = + get_file_io_with_props(vec![(S3_MULTIPART_PART_SIZE_BYTES, PART_SIZE.to_string())]) + .await; + let key = normalize_test_name_with_parts!( + "test_file_io_s3_writer_splits_write_into_bounded_parts" + ); + let path = format!("s3://{OBJECT_STORE_BUCKET}/{key}"); + + let _ = file_io.delete(&path).await; + let mut writer = file_io.new_output(&path).unwrap().writer().await.unwrap(); + writer + .write(Bytes::from(vec![0u8; PART_SIZE * PARTS])) + .await + .unwrap(); + writer.close().await.unwrap(); + + assert_eq!(upload_part_count(&key).await, PARTS); + file_io.delete(&path).await.unwrap(); + } + + #[tokio::test] + async fn test_file_io_s3_write_splits_buffer_into_bounded_parts() { + const PART_SIZE: usize = 5 * 1024 * 1024; + const PARTS: usize = 3; + + let file_io = + get_file_io_with_props(vec![(S3_MULTIPART_PART_SIZE_BYTES, PART_SIZE.to_string())]) + .await; + let key = normalize_test_name_with_parts!( + "test_file_io_s3_write_splits_buffer_into_bounded_parts" + ); + let path = format!("s3://{OBJECT_STORE_BUCKET}/{key}"); + + let _ = file_io.delete(&path).await; + file_io + .new_output(&path) + .unwrap() + .write(Bytes::from(vec![0u8; PART_SIZE * PARTS])) + .await + .unwrap(); + + assert_eq!(upload_part_count(&key).await, PARTS); + file_io.delete(&path).await.unwrap(); + } + + #[tokio::test] + async fn test_file_io_s3_write_below_part_size_stays_single_request() { + let file_io = get_file_io_with_props(vec![( + S3_MULTIPART_PART_SIZE_BYTES, + (5 * 1024 * 1024).to_string(), + )]) + .await; + let key = normalize_test_name_with_parts!( + "test_file_io_s3_write_below_part_size_stays_single_request" + ); + let path = format!("s3://{OBJECT_STORE_BUCKET}/{key}"); + + let _ = file_io.delete(&path).await; + file_io + .new_output(&path) + .unwrap() + .write(Bytes::from_static(b"{\"format-version\":2}")) + .await + .unwrap(); + + assert_eq!(upload_part_count(&key).await, 1); + file_io.delete(&path).await.unwrap(); + } + #[tokio::test] async fn test_file_io_s3_input() { let file_io = get_file_io().await; // Use unique file path based on module path to avoid conflicts let file_path = format!( - "s3://bucket1/{}", + "s3://{OBJECT_STORE_BUCKET}/{}", normalize_test_name_with_parts!("test_file_io_s3_input") ); let output_file = file_io.new_output(&file_path).unwrap(); @@ -199,7 +327,10 @@ mod tests { .build(); // Test that the FileIO was built successfully with the custom loader - match file_io_with_custom_creds.exists("s3://bucket1/any").await { + match file_io_with_custom_creds + .exists(&format!("s3://{OBJECT_STORE_BUCKET}/any")) + .await + { Ok(_) => {} Err(e) => panic!("Failed to check existence of bucket: {e}"), } @@ -227,7 +358,10 @@ mod tests { .build(); // Test that the FileIO was built successfully with the custom loader - match file_io_with_custom_creds.exists("s3://bucket1/any").await { + match file_io_with_custom_creds + .exists(&format!("s3://{OBJECT_STORE_BUCKET}/any")) + .await + { Ok(_) => panic!( "Expected error, but got Ok - the credential loader should fail to provide valid credentials" ), @@ -248,7 +382,7 @@ mod tests { let paths: Vec = (0..5) .map(|i| { format!( - "s3://bucket1/{}/file-{i}", + "s3://{OBJECT_STORE_BUCKET}/{}/file-{i}", normalize_test_name_with_parts!("test_file_io_s3_delete_stream") ) })