This is an automated email from the ASF dual-hosted git repository. hubcio pushed a commit to branch fix/iceberg-partition-routing in repository https://gitbox.apache.org/repos/asf/iggy.git
commit 14649bae09c429ca6d4fca88251c72d209dd3ede Author: Hubert Gruszecki <[email protected]> AuthorDate: Wed Aug 26 10:28:13 2026 +0200 fix(connectors): write iceberg sink data into real table partitions The Iceberg sink built its partition key from an empty PartitionSpec::builder() with hardcoded zero literals, so every record landed in one dummy partition (e.g. year=1970) whatever the table's spec said, and a batch spanning several partitions was never split. Read the default partition spec from table metadata instead and push every batch through one TableWriter: partitioned tables split it with RecordBatchPartitionSplitter and fan out to one data file writer per partition value, unpartitioned tables map it to the single key of their spec (with one null per void field, which the commit-time arity check requires). Files finalized by a failed batch are deleted before the error is returned so nothing is left orphaned on the object store; files are kept on commit failure because the catalog may already have applied the commit. Ports #3860 with its review feedback applied. Fixes #3853 --- Cargo.lock | 2 + core/connectors/sinks/iceberg_sink/Cargo.toml | 4 + core/connectors/sinks/iceberg_sink/README.md | 7 + .../sinks/iceberg_sink/src/router/mod.rs | 413 +++++++++++++++------ .../tests/connectors/fixtures/iceberg/container.rs | 64 +++- .../tests/connectors/fixtures/iceberg/mod.rs | 3 +- core/integration/tests/connectors/fixtures/mod.rs | 3 +- .../tests/connectors/iceberg/iceberg_sink.rs | 82 +++- 8 files changed, 465 insertions(+), 113 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 251f7f309..1eb6e89f3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7053,6 +7053,7 @@ dependencies = [ name = "iggy_connector_iceberg_sink" version = "0.5.0-edge.4" dependencies = [ + "arrow-array 57.3.1", "arrow-json 57.3.1", "async-trait", "dashmap", @@ -7064,6 +7065,7 @@ dependencies = [ "serde", "simd-json", "strum 0.28.0", + "tokio", "tracing", "uuid", ] diff --git a/core/connectors/sinks/iceberg_sink/Cargo.toml b/core/connectors/sinks/iceberg_sink/Cargo.toml index ec6574214..9e5e21cd9 100644 --- a/core/connectors/sinks/iceberg_sink/Cargo.toml +++ b/core/connectors/sinks/iceberg_sink/Cargo.toml @@ -35,6 +35,7 @@ ignored = ["dashmap"] crate-type = ["cdylib", "lib"] [dependencies] +arrow-array = { workspace = true } arrow-json = { workspace = true } async-trait = { workspace = true } dashmap = { workspace = true } @@ -49,5 +50,8 @@ strum = { workspace = true } tracing = { workspace = true } uuid = { workspace = true } +[dev-dependencies] +tokio = { workspace = true } + [lints] workspace = true diff --git a/core/connectors/sinks/iceberg_sink/README.md b/core/connectors/sinks/iceberg_sink/README.md index c9079e713..2ebb4405f 100644 --- a/core/connectors/sinks/iceberg_sink/README.md +++ b/core/connectors/sinks/iceberg_sink/README.md @@ -81,6 +81,13 @@ Example: - Namespace: `nyc` - Table name: `users` +## Partitioned Tables + +Data files follow the table's default partition spec. Each batch is split by the spec's transforms +(`identity`, `bucket`, `truncate`, `year`, `month`, `day`, `hour`) and every partition value present +in the batch gets at least one Parquet file under its partition path. Unpartitioned tables are +written into a single file per batch. + ## Source Compatibility The Iceberg sink expects **flat JSON** where each top-level key maps directly to a column in the diff --git a/core/connectors/sinks/iceberg_sink/src/router/mod.rs b/core/connectors/sinks/iceberg_sink/src/router/mod.rs index 3f43b9401..39fa7197e 100644 --- a/core/connectors/sinks/iceberg_sink/src/router/mod.rs +++ b/core/connectors/sinks/iceberg_sink/src/router/mod.rs @@ -17,16 +17,16 @@ use crate::router::arrow_streamer::JsonArrowReader; use crate::slice_user_table; +use arrow_array::RecordBatch; use arrow_json::ReaderBuilder; use async_trait::async_trait; use iceberg::TableIdent; -use iceberg::arrow::schema_to_arrow_schema; -use iceberg::spec::{ - Literal, PartitionKey, PartitionSpec, PrimitiveLiteral, PrimitiveType, Struct, StructType, -}; +use iceberg::arrow::{RecordBatchPartitionSplitter, schema_to_arrow_schema}; +use iceberg::io::FileIO; +use iceberg::spec::{DataFile, PartitionKey, Struct}; use iceberg::table::Table; use iceberg::transaction::{ApplyTransactionAction, Transaction}; -use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder; +use iceberg::writer::base_writer::data_file_writer::{DataFileWriter, DataFileWriterBuilder}; use iceberg::writer::file_writer::ParquetWriterBuilder; use iceberg::writer::file_writer::rolling_writer::RollingFileWriterBuilder; use iceberg::writer::{IcebergWriter, IcebergWriterBuilder}; @@ -36,10 +36,17 @@ use iceberg::{ }; use iggy_connector_sdk::{ConsumedMessage, Error, MessagesMetadata, Payload, Schema}; use parquet::file::properties::WriterProperties; +use std::collections::HashMap; +use std::collections::hash_map::Entry; use std::sync::Arc; use tracing::{error, warn}; use uuid::Uuid; +type ParquetDataFileWriterBuilder = + DataFileWriterBuilder<ParquetWriterBuilder, DefaultLocationGenerator, DefaultFileNameGenerator>; +type ParquetDataFileWriter = + DataFileWriter<ParquetWriterBuilder, DefaultLocationGenerator, DefaultFileNameGenerator>; + fn format_error_chain(err: &dyn std::error::Error) -> String { let mut chain = err.to_string(); let mut source = err.source(); @@ -67,55 +74,49 @@ async fn table_exists(route_field_val: &str, catalog: &dyn Catalog) -> Option<Ta catalog.load_table(&table_ident).await.ok() } -pub fn primitive_type_to_literal(pt: &PrimitiveType) -> Result<PrimitiveLiteral, Error> { - match pt { - PrimitiveType::Boolean => Ok(PrimitiveLiteral::Boolean(false)), - PrimitiveType::Int => Ok(PrimitiveLiteral::Int(0)), - PrimitiveType::Long => Ok(PrimitiveLiteral::Long(0)), - PrimitiveType::Decimal { .. } => Ok(PrimitiveLiteral::Int128(0)), - PrimitiveType::Date => Ok(PrimitiveLiteral::Int(0)), // e.g. days since epoch - PrimitiveType::Time => Ok(PrimitiveLiteral::Long(0)), // microseconds since midnight - PrimitiveType::Timestamp => Ok(PrimitiveLiteral::Long(0)), // microseconds since epoch - PrimitiveType::Timestamptz => Ok(PrimitiveLiteral::Long(0)), - PrimitiveType::TimestampNs => Ok(PrimitiveLiteral::Long(0)), - PrimitiveType::TimestamptzNs => Ok(PrimitiveLiteral::Long(0)), - PrimitiveType::String => Ok(PrimitiveLiteral::String(String::new())), - PrimitiveType::Uuid => Ok(PrimitiveLiteral::Binary(vec![0; 16])), - PrimitiveType::Fixed(len) => Ok(PrimitiveLiteral::Binary(vec![0; *len as usize])), - PrimitiveType::Binary => Ok(PrimitiveLiteral::Binary(Vec::new())), - _ => { - error!("Partition type not supported"); - Err(Error::InvalidConfig) - } - } -} - -fn get_partition_type_value(default_partition_type: &StructType) -> Result<Option<Struct>, Error> { - let mut fields: Vec<Option<Literal>> = Vec::new(); +async fn write_data( + messages: &[Payload], + table: &Table, + catalog: &dyn Catalog, + messages_schema: Schema, +) -> Result<(), Error> { + let data_files = write_data_files(messages, table, messages_schema).await?; - if default_partition_type.fields().is_empty() { - return Ok(None); - }; + // Data files are kept on commit failure: the catalog may have applied the + // commit before the error surfaced, and retries must check that first. + let table_commit = Transaction::new(table); - for field in default_partition_type.fields() { - let field_type = field.field_type.as_primitive_type().ok_or_else(|| { - error!("The partition type of the configured iceberg table is not a primitive type"); - Error::InvalidConfig - })?; + let action = table_commit.fast_append().add_data_files(data_files); - let value = Some(Literal::Primitive(primitive_type_to_literal(field_type)?)); + let tx = action.apply(table_commit).map_err(|err| { + let chain = format_error_chain(&err); + error!( + "Failed to apply transaction on table with UUID: {}, Error: {}", + table.metadata().uuid(), + chain + ); + Error::TransactionApplyError(chain) + })?; - fields.push(value); - } - Ok(Some(Struct::from_iter(fields))) + let _table = tx.commit(catalog).await.map_err(|err| { + let chain = format_error_chain(&err); + error!( + "Failed to commit transaction on table with UUID: {}, Error: {}", + table.metadata().uuid(), + chain + ); + Error::CatalogCommitError(chain) + })?; + Ok(()) } -async fn write_data( +/// Writes the JSON payloads as Parquet data files under the table location, +/// laid out by the table's default partition spec. Nothing is committed here. +async fn write_data_files( messages: &[Payload], table: &Table, - catalog: &dyn Catalog, messages_schema: Schema, -) -> Result<(), Error> { +) -> Result<Vec<DataFile>, Error> { let location = DefaultLocationGenerator::new(table.metadata().clone()).map_err(|err| { error!( "Failed to get location on table: {}. Error: {}", @@ -145,27 +146,6 @@ async fn write_data( let data_file_writer_builder = DataFileWriterBuilder::new(rolling_file_writer_builder); - let partition_spec = PartitionSpec::builder(table.current_schema_ref()); - - let partition_type = get_partition_type_value(table.metadata().default_partition_type())?; - - let mut writer = data_file_writer_builder - .build(match partition_type { - None => None, - Some(p_type) => Some(PartitionKey::new( - partition_spec - .build() - .map_err(|err| Error::InitError(err.to_string()))?, - table.current_schema_ref(), - p_type, - )), - }) - .await - .map_err(|err| { - error!("Error while constructing data file writer: {}", err); - Error::InitError(err.to_string()) - })?; - let msgs: Vec<&simd_json::OwnedValue> = messages .iter() .filter_map(|payload| match payload { @@ -201,6 +181,12 @@ async fn write_data( Error::InitError(err.to_string()) })?; + let mut writer = TableWriter::new(table, data_file_writer_builder).map_err(|err| { + let chain = format_error_chain(&err); + error!("Error while constructing data file writer: {}", chain); + Error::InitError(chain) + })?; + let write_result: Result<(), Error> = async { for batch in reader { let batch_data = batch.map_err(|err| { @@ -208,60 +194,138 @@ async fn write_data( error!("Error while getting record batch: {}", chain); Error::InvalidRecordValue(chain) })?; - writer.write(batch_data).await.map_err(|err| { - let chain = format_error_chain(&err); - error!("Error while writing record batch: {}", chain); - Error::WriteFailure(chain) - })?; + writer.write(batch_data).await?; } Ok(()) } .await; - if let Err(e) = &write_result { + let (data_files, close_error) = writer.close().await; + if let Some(err) = write_result.err().or(close_error) { error!( - "Batch loop failed ({}), closing writer to release resources", - e + "Writing data files failed ({}), deleting the uncommitted ones", + err ); - if let Err(close_err) = writer.close().await { - error!("Failed to close writer after batch error: {}", close_err); + delete_uncommitted_files(table.file_io(), &data_files).await; + return Err(err); + } + Ok(data_files) +} + +/// Files finalized by a failed batch are never committed, so they are removed +/// instead of being left as orphans on the object store. +async fn delete_uncommitted_files(file_io: &FileIO, data_files: &[DataFile]) { + for data_file in data_files { + if let Err(err) = file_io.delete(data_file.file_path()).await { + warn!( + "Failed to delete uncommitted data file {}: {}", + data_file.file_path(), + err + ); } - return Err(write_result.unwrap_err()); } +} - let data_files = writer.close().await.map_err(|err| { - let chain = format_error_chain(&err); - error!( - "Error while writing data records to Parquet file: {}", - chain - ); - Error::WriteFailure(chain) - })?; +/// Data file writer bound to the table's default partition spec. +/// +/// Every batch is mapped to partition keys and fanned out to one data file +/// writer per key, so a mixed batch lands in the right partitions. An +/// unpartitioned table maps every batch to the single key of its spec. +/// +/// The per-partition writers are owned here rather than by iceberg's +/// `FanoutWriter` because its `close` stops at the first failing writer and +/// loses the files the other writers already finalized. +struct TableWriter { + partitioner: Partitioner, + builder: ParquetDataFileWriterBuilder, + writers: HashMap<Struct, ParquetDataFileWriter>, +} - let table_commit = Transaction::new(table); +enum Partitioner { + Unpartitioned(PartitionKey), + /// Boxed to keep the enum small (`clippy::large_enum_variant`). + Partitioned(Box<RecordBatchPartitionSplitter>), +} - let action = table_commit.fast_append().add_data_files(data_files); +impl TableWriter { + fn new(table: &Table, builder: ParquetDataFileWriterBuilder) -> iceberg::Result<Self> { + let partition_spec = table.metadata().default_partition_spec(); + let schema = table.metadata().current_schema().clone(); + let partitioner = if partition_spec.is_unpartitioned() { + // A spec whose fields are all `Transform::Void` counts as unpartitioned, + // but its partition type still has one field per void field and the + // commit rejects data files whose partition value has a different arity. + let void_field_count = table.metadata().default_partition_type().fields().len(); + Partitioner::Unpartitioned(PartitionKey::new( + partition_spec.as_ref().clone(), + schema, + Struct::from_iter(vec![None; void_field_count]), + )) + } else { + let splitter = RecordBatchPartitionSplitter::try_new_with_computed_values( + schema, + partition_spec.clone(), + )?; + Partitioner::Partitioned(Box::new(splitter)) + }; + Ok(Self { + partitioner, + builder, + writers: HashMap::new(), + }) + } - let tx = action.apply(table_commit).map_err(|err| { - let chain = format_error_chain(&err); - error!( - "Failed to apply transaction on table with UUID: {}, Error: {}", - table.metadata().uuid(), - chain - ); - Error::TransactionApplyError(chain) - })?; + async fn write(&mut self, batch: RecordBatch) -> Result<(), Error> { + let partitioned_batches = match &self.partitioner { + Partitioner::Unpartitioned(partition_key) => vec![(partition_key.clone(), batch)], + Partitioner::Partitioned(splitter) => splitter.split(&batch).map_err(|err| { + let chain = format_error_chain(&err); + error!("Error while splitting record batch by partition: {}", chain); + Error::InvalidRecordValue(chain) + })?, + }; + for (partition_key, partition_batch) in partitioned_batches { + let writer = match self.writers.entry(partition_key.data().clone()) { + Entry::Occupied(entry) => entry.into_mut(), + Entry::Vacant(entry) => { + let writer = self + .builder + .build(Some(partition_key)) + .await + .map_err(write_failure)?; + entry.insert(writer) + } + }; + writer.write(partition_batch).await.map_err(write_failure)?; + } + Ok(()) + } - let _table = tx.commit(catalog).await.map_err(|err| { - let chain = format_error_chain(&err); - error!( - "Failed to commit transaction on table with UUID: {}, Error: {}", - table.metadata().uuid(), - chain - ); - Error::CatalogCommitError(chain) - })?; - Ok(()) + /// Closes every partition writer even after one of them fails, so the + /// caller learns about all finalized files and can commit or delete them. + /// Returns the first close failure alongside the files. + async fn close(self) -> (Vec<DataFile>, Option<Error>) { + let mut data_files = Vec::new(); + let mut first_error = None; + for mut writer in self.writers.into_values() { + match writer.close().await { + Ok(files) => data_files.extend(files), + Err(err) => { + let close_error = write_failure(err); + if first_error.is_none() { + first_error = Some(close_error); + } + } + } + } + (data_files, first_error) + } +} + +fn write_failure(err: iceberg::Error) -> Error { + let chain = format_error_chain(&err); + error!("Error while writing data file: {}", chain); + Error::WriteFailure(chain) } #[async_trait] @@ -272,3 +336,140 @@ pub trait Router: std::fmt::Debug + Sync + Send { messages: Vec<ConsumedMessage>, ) -> Result<(), crate::Error>; } + +#[cfg(test)] +mod tests { + use super::*; + use iceberg::spec::{ + FormatVersion, Literal, NestedField, PrimitiveLiteral, PrimitiveType, + Schema as IcebergSchema, SortOrder, TableMetadataBuilder, Transform, Type, + UnboundPartitionSpec, + }; + use std::collections::HashMap; + + const REGION_FIELD_ID: i32 = 2; + + fn in_memory_table(partition_spec: UnboundPartitionSpec) -> Table { + let schema = IcebergSchema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(), + NestedField::required( + REGION_FIELD_ID, + "region", + Type::Primitive(PrimitiveType::String), + ) + .into(), + ]) + .build() + .expect("Failed to build schema"); + let metadata = TableMetadataBuilder::new( + schema, + partition_spec, + SortOrder::unsorted_order(), + "memory://warehouse/test/events".to_string(), + FormatVersion::V2, + HashMap::new(), + ) + .expect("Failed to create table metadata") + .build() + .expect("Failed to build table metadata") + .metadata; + + Table::builder() + .identifier(TableIdent::from_strs(["test", "events"]).expect("Failed to build ident")) + .metadata(metadata) + .file_io(FileIO::new_with_memory()) + .build() + .expect("Failed to build table") + } + + fn json_payloads(rows: &[(i64, &str)]) -> Vec<Payload> { + rows.iter() + .map(|(id, region)| Payload::Json(simd_json::json!({ "id": *id, "region": *region }))) + .collect() + } + + fn write(table: &Table, payloads: &[Payload]) -> Vec<DataFile> { + tokio::runtime::Runtime::new() + .expect("Failed to create runtime") + .block_on(write_data_files(payloads, table, Schema::Json)) + .expect("Failed to write data files") + } + + fn partition_region(data_file: &DataFile) -> String { + match &data_file.partition()[0] { + Some(Literal::Primitive(PrimitiveLiteral::String(region))) => region.clone(), + other => panic!("Expected string partition value, got {other:?}"), + } + } + + #[test] + fn given_partitioned_table_should_write_one_data_file_per_partition_value() { + let partition_spec = UnboundPartitionSpec::builder() + .add_partition_field(REGION_FIELD_ID, "region", Transform::Identity) + .expect("Failed to add partition field") + .build(); + let table = in_memory_table(partition_spec); + let payloads = json_payloads(&[(1, "eu"), (2, "us"), (3, "eu")]); + + let data_files = write(&table, &payloads); + + assert_eq!(data_files.len(), 2); + let by_region: HashMap<String, &DataFile> = data_files + .iter() + .map(|data_file| (partition_region(data_file), data_file)) + .collect(); + assert_eq!(by_region["eu"].record_count(), 2); + assert!(by_region["eu"].file_path().contains("/data/region=eu/")); + assert_eq!(by_region["us"].record_count(), 1); + assert!(by_region["us"].file_path().contains("/data/region=us/")); + } + + #[test] + fn given_void_only_partition_spec_should_write_null_partition_values() { + let partition_spec = UnboundPartitionSpec::builder() + .add_partition_field(REGION_FIELD_ID, "region_void", Transform::Void) + .expect("Failed to add partition field") + .build(); + let table = in_memory_table(partition_spec); + let payloads = json_payloads(&[(1, "eu"), (2, "us")]); + + let data_files = write(&table, &payloads); + + assert_eq!(data_files.len(), 1); + assert_eq!(data_files[0].partition().fields(), &[None]); + assert!(!data_files[0].file_path().contains("region_void=")); + } + + #[test] + fn given_uncommitted_files_should_be_deleted_from_store() { + let table = in_memory_table(UnboundPartitionSpec::builder().build()); + let payloads = json_payloads(&[(1, "eu")]); + let data_files = write(&table, &payloads); + let file_path = data_files[0].file_path().to_string(); + + let exists_after_delete = tokio::runtime::Runtime::new() + .expect("Failed to create runtime") + .block_on(async { + assert!(table.file_io().exists(&file_path).await.expect("exists")); + delete_uncommitted_files(table.file_io(), &data_files).await; + table.file_io().exists(&file_path).await.expect("exists") + }); + + assert!(!exists_after_delete); + } + + #[test] + fn given_unpartitioned_table_should_write_single_data_file_without_partition() { + let table = in_memory_table(UnboundPartitionSpec::builder().build()); + let payloads = json_payloads(&[(1, "eu"), (2, "us")]); + + let data_files = write(&table, &payloads); + + assert_eq!(data_files.len(), 1); + assert_eq!(data_files[0].record_count(), 2); + assert!(data_files[0].partition().fields().is_empty()); + assert!(data_files[0].file_path().contains("/data/")); + assert!(!data_files[0].file_path().contains("region=")); + } +} diff --git a/core/integration/tests/connectors/fixtures/iceberg/container.rs b/core/integration/tests/connectors/fixtures/iceberg/container.rs index cf5386c8b..71332f0ab 100644 --- a/core/integration/tests/connectors/fixtures/iceberg/container.rs +++ b/core/integration/tests/connectors/fixtures/iceberg/container.rs @@ -275,10 +275,11 @@ pub trait IcebergOps: Sync { &self, namespace: &str, table: &str, + partition_spec: Option<serde_json::Value>, ) -> impl std::future::Future<Output = Result<(), TestBinaryError>> + Send { async move { let url = format!("{}/v1/namespaces/{namespace}/tables", self.catalog_url()); - let body = serde_json::json!({ + let mut body = serde_json::json!({ "name": table, "schema": { "type": "struct", @@ -292,6 +293,9 @@ pub trait IcebergOps: Sync { ] } }); + if let Some(partition_spec) = partition_spec { + body["partition-spec"] = partition_spec; + } let response = self .http_client() @@ -431,6 +435,8 @@ pub struct SnapshotSummary { pub added_data_files: Option<String>, #[serde(rename = "added-records")] pub added_records: Option<String>, + #[serde(rename = "changed-partition-count")] + pub changed_partition_count: Option<String>, } impl IcebergOps for IcebergFixture { @@ -502,6 +508,8 @@ fn create_http_client() -> HttpClient { pub const DEFAULT_NAMESPACE: &str = "test"; pub const DEFAULT_TABLE: &str = "messages"; +/// Field id of the `active` column in the schema created by `IcebergOps::create_table`. +const ACTIVE_FIELD_ID: i32 = 5; pub struct IcebergPreCreatedFixture { inner: IcebergFixture, @@ -524,13 +532,61 @@ impl IcebergOps for IcebergPreCreatedFixture { } } +impl IcebergPreCreatedFixture { + async fn setup_with_partition_spec( + partition_spec: Option<serde_json::Value>, + ) -> Result<Self, TestBinaryError> { + let inner = IcebergFixture::setup().await?; + + inner.create_namespace(DEFAULT_NAMESPACE).await?; + inner + .create_table(DEFAULT_NAMESPACE, DEFAULT_TABLE, partition_spec) + .await?; + + Ok(Self { inner }) + } +} + #[async_trait] impl TestFixture for IcebergPreCreatedFixture { async fn setup() -> Result<Self, TestBinaryError> { - let inner = IcebergFixture::setup().await?; + Self::setup_with_partition_spec(None).await + } - inner.create_namespace(DEFAULT_NAMESPACE).await?; - inner.create_table(DEFAULT_NAMESPACE, DEFAULT_TABLE).await?; + fn connectors_runtime_envs(&self) -> HashMap<String, String> { + self.inner.connectors_runtime_envs() + } +} + +/// Same table as `IcebergPreCreatedFixture`, partitioned by identity(`active`). +pub struct IcebergPartitionedTableFixture { + inner: IcebergPreCreatedFixture, +} + +impl IcebergOps for IcebergPartitionedTableFixture { + fn catalog_url(&self) -> &str { + self.inner.catalog_url() + } + + fn http_client(&self) -> &HttpClient { + self.inner.http_client() + } +} + +#[async_trait] +impl TestFixture for IcebergPartitionedTableFixture { + async fn setup() -> Result<Self, TestBinaryError> { + let partition_spec = serde_json::json!({ + "spec-id": 0, + "fields": [{ + "source-id": ACTIVE_FIELD_ID, + "field-id": 1000, + "name": "active", + "transform": "identity" + }] + }); + let inner = + IcebergPreCreatedFixture::setup_with_partition_spec(Some(partition_spec)).await?; Ok(Self { inner }) } diff --git a/core/integration/tests/connectors/fixtures/iceberg/mod.rs b/core/integration/tests/connectors/fixtures/iceberg/mod.rs index 670c2513c..0b5dd1877 100644 --- a/core/integration/tests/connectors/fixtures/iceberg/mod.rs +++ b/core/integration/tests/connectors/fixtures/iceberg/mod.rs @@ -18,5 +18,6 @@ mod container; pub use container::{ - DEFAULT_NAMESPACE, DEFAULT_TABLE, IcebergEnvAuthFixture, IcebergOps, IcebergPreCreatedFixture, + DEFAULT_NAMESPACE, DEFAULT_TABLE, IcebergEnvAuthFixture, IcebergOps, + IcebergPartitionedTableFixture, IcebergPreCreatedFixture, }; diff --git a/core/integration/tests/connectors/fixtures/mod.rs b/core/integration/tests/connectors/fixtures/mod.rs index 7eaf6fe51..e9d7a1332 100644 --- a/core/integration/tests/connectors/fixtures/mod.rs +++ b/core/integration/tests/connectors/fixtures/mod.rs @@ -63,7 +63,8 @@ pub use http::{ HttpSinkNdjsonFixture, HttpSinkNoMetadataFixture, HttpSinkRawFixture, }; pub use iceberg::{ - DEFAULT_NAMESPACE, DEFAULT_TABLE, IcebergEnvAuthFixture, IcebergOps, IcebergPreCreatedFixture, + DEFAULT_NAMESPACE, DEFAULT_TABLE, IcebergEnvAuthFixture, IcebergOps, + IcebergPartitionedTableFixture, IcebergPreCreatedFixture, }; pub use influxdb::{ InfluxDb3SinkFixture, InfluxDb3SourceFixture, InfluxDbSinkBase64Fixture, InfluxDbSinkFixture, diff --git a/core/integration/tests/connectors/iceberg/iceberg_sink.rs b/core/integration/tests/connectors/iceberg/iceberg_sink.rs index ceb082f6f..2e969775c 100644 --- a/core/integration/tests/connectors/iceberg/iceberg_sink.rs +++ b/core/integration/tests/connectors/iceberg/iceberg_sink.rs @@ -17,7 +17,8 @@ use crate::connectors::create_test_messages; use crate::connectors::fixtures::{ - DEFAULT_NAMESPACE, DEFAULT_TABLE, IcebergEnvAuthFixture, IcebergOps, IcebergPreCreatedFixture, + DEFAULT_NAMESPACE, DEFAULT_TABLE, IcebergEnvAuthFixture, IcebergOps, + IcebergPartitionedTableFixture, IcebergPreCreatedFixture, }; use bytes::Bytes; use iggy::prelude::{IggyMessage, Partitioning}; @@ -33,6 +34,9 @@ const ICEBERG_SINK_KEY: &str = "iceberg"; const SNAPSHOT_POLL_ATTEMPTS: usize = 30; const SNAPSHOT_POLL_INTERVAL_MS: u64 = 500; const BULK_SNAPSHOT_POLL_ATTEMPTS: usize = 60; +/// `create_test_messages` alternates `active`, so this many messages span both partition values. +const PARTITIONED_MESSAGE_COUNT: usize = 5; +const ACTIVE_PARTITION_COUNT: usize = 2; #[iggy_harness( server(connectors_runtime(config_path = "tests/connectors/iceberg/sink.toml")), @@ -206,6 +210,82 @@ async fn iceberg_sink_handles_bulk_messages( assert!(sinks[0].last_error.is_none()); } +#[iggy_harness( + server(connectors_runtime(config_path = "tests/connectors/iceberg/sink.toml")), + seed = seeds::connector_stream +)] +async fn iceberg_sink_routes_messages_to_table_partitions( + harness: &TestHarness, + fixture: IcebergPartitionedTableFixture, +) { + let client = harness.root_client().await.unwrap(); + + let stream_id: Identifier = seeds::names::STREAM.try_into().unwrap(); + let topic_id: Identifier = seeds::names::TOPIC.try_into().unwrap(); + + let test_messages = create_test_messages(PARTITIONED_MESSAGE_COUNT); + let mut messages: Vec<IggyMessage> = test_messages + .iter() + .enumerate() + .map(|(i, msg)| { + let payload = serde_json::to_vec(msg).expect("Failed to serialize message"); + IggyMessage::builder() + .id((i + 1) as u128) + .payload(Bytes::from(payload)) + .build() + .expect("Failed to build message") + }) + .collect(); + + client + .send_messages( + &stream_id, + &topic_id, + &Partitioning::partition_id(0), + &mut messages, + ) + .await + .expect("Failed to send messages"); + + fixture + .wait_for_snapshots( + DEFAULT_NAMESPACE, + DEFAULT_TABLE, + 1, + SNAPSHOT_POLL_ATTEMPTS, + SNAPSHOT_POLL_INTERVAL_MS, + ) + .await + .expect("Data should be written to Iceberg table"); + + let metadata = fixture + .get_table_metadata(DEFAULT_NAMESPACE, DEFAULT_TABLE) + .await + .expect("Failed to get table metadata"); + let snapshots = metadata + .metadata + .snapshots + .expect("Table should have a snapshot"); + assert_eq!(snapshots.len(), 1, "One batch must produce one snapshot"); + let summary = snapshots[0] + .summary + .as_ref() + .expect("Snapshot should have a summary"); + + assert_eq!( + summary.added_records.as_deref(), + Some(PARTITIONED_MESSAGE_COUNT.to_string().as_str()) + ); + assert_eq!( + summary.changed_partition_count.as_deref(), + Some(ACTIVE_PARTITION_COUNT.to_string().as_str()) + ); + assert_eq!( + summary.added_data_files.as_deref(), + Some(ACTIVE_PARTITION_COUNT.to_string().as_str()) + ); +} + #[iggy_harness( server(connectors_runtime(config_path = "tests/connectors/iceberg/sink_default_credentials.toml")), seed = seeds::connector_stream
