From fe678adeb45133e6309dacad519de93c03ee96c5 Mon Sep 17 00:00:00 2001 From: Sergei Turukin Date: Wed, 15 Apr 2026 15:13:26 -0500 Subject: [PATCH 1/6] fix: v2 manifest-list field renames + TIMESTAMPTZ date_transform + identity partition column filter Three independent fixes needed to read and scan real v2 Iceberg tables written by current Apache Iceberg (>= 1.0) and partitioned by an identity column or a day/hour transform on TIMESTAMPTZ. 1. `iceberg-rust-spec/src/spec/manifest_list.rs` - the v2 manifest_list Avro schema uses `added_data_files_count` / `existing_data_files_count` / `deleted_data_files_count`, but the reader still used the older `added_files_count` / `existing_files_count` / `deleted_files_count` names. Any manifest list written by modern Apache Iceberg failed to deserialize with "field not found" before the reader even reached an entry. Declare each count with `#[serde(rename = "added_data_files_count", alias = "added_files_count")]` so both new and legacy field names resolve cleanly, and update the static reader Avro schema to emit the current names. New regression test `test_manifest_list_v2_apache_field_names` simulates an Apache Iceberg >= 1.0 writer and asserts the reader deserializes it. 2. `datafusion_iceberg/src/pruning_statistics.rs` - the internal `DateTransform` UDF used a hardcoded `OneOf(Exact([Utf8, Date32]), Exact([Utf8, Timestamp(us, None)]))` signature, so any timezone-aware timestamp fell through with a type-check error. Replace with `TypeSignature::UserDefined` plus a `coerce_types` impl that accepts any `Timestamp(*, *)` and normalizes to `Timestamp(Microsecond, None)` for the physical call. Partition transforms operate on i64 microseconds-since-epoch and are timezone- agnostic, so stripping the tz on input is safe. 3. `datafusion_iceberg/src/table.rs::datafusion_partition_columns` - skip partition fields whose transform is `Identity` and whose name equals the source column name. For those, Iceberg materializes the column both in the parquet file body and in the Hive-style directory encoding; DataFusion's parquet reader then trips on an off-by-one ("expected N cols but got N+1") because it tries to derive the same column from both places. A follow-up commit promotes this filter out of `datafusion_partition_columns` so the manifest pruner sees the same filtered list. --- datafusion_iceberg/src/pruning_statistics.rs | 46 ++++++-- datafusion_iceberg/src/table.rs | 13 +++ iceberg-rust-spec/src/spec/manifest_list.rs | 106 ++++++++++++++++++- 3 files changed, 155 insertions(+), 10 deletions(-) diff --git a/datafusion_iceberg/src/pruning_statistics.rs b/datafusion_iceberg/src/pruning_statistics.rs index c1be1947..85ebc4c8 100644 --- a/datafusion_iceberg/src/pruning_statistics.rs +++ b/datafusion_iceberg/src/pruning_statistics.rs @@ -357,14 +357,14 @@ struct DateTransform { impl DateTransform { fn new() -> Self { + // Accept any second-argument type via `TypeSignature::UserDefined` and + // normalize it in `coerce_types`. The underlying transform is + // timezone-agnostic (it operates on the i64 microseconds-since-epoch), + // so any `Timestamp(Microsecond, *)` is a valid input — we just need + // to strip the timezone metadata so the physical invocation sees + // `Timestamp(Microsecond, None)`. let signature = Signature { - type_signature: TypeSignature::OneOf(vec![ - TypeSignature::Exact(vec![DataType::Utf8, DataType::Date32]), - TypeSignature::Exact(vec![ - DataType::Utf8, - DataType::Timestamp(datafusion::arrow::datatypes::TimeUnit::Microsecond, None), - ]), - ]), + type_signature: TypeSignature::UserDefined, volatility: Volatility::Immutable, }; Self { signature } @@ -388,6 +388,38 @@ impl ScalarUDFImpl for DateTransform { Ok(DataType::Int32) } + fn coerce_types( + &self, + arg_types: &[DataType], + ) -> datafusion::error::Result> { + use datafusion::arrow::datatypes::TimeUnit; + if arg_types.len() != 2 { + return Err(DataFusionError::Plan(format!( + "date_transform expects 2 arguments, got {}", + arg_types.len() + ))); + } + if !matches!(arg_types[0], DataType::Utf8 | DataType::LargeUtf8) { + return Err(DataFusionError::Plan(format!( + "date_transform first argument must be Utf8, got {}", + arg_types[0] + ))); + } + let coerced_second = match &arg_types[1] { + DataType::Date32 => DataType::Date32, + DataType::Timestamp(TimeUnit::Microsecond, _) => { + DataType::Timestamp(TimeUnit::Microsecond, None) + } + DataType::Timestamp(unit, _) => DataType::Timestamp(*unit, None), + other => { + return Err(DataFusionError::Plan(format!( + "date_transform second argument must be Date32 or Timestamp, got {other}" + ))) + } + }; + Ok(vec![DataType::Utf8, coerced_second]) + } + fn invoke_with_args( &self, args: ScalarFunctionArgs, diff --git a/datafusion_iceberg/src/table.rs b/datafusion_iceberg/src/table.rs index ee5c64bc..67d989bf 100644 --- a/datafusion_iceberg/src/table.rs +++ b/datafusion_iceberg/src/table.rs @@ -932,8 +932,21 @@ async fn table_scan( fn datafusion_partition_columns( partition_fields: &[BoundPartitionField<'_>], ) -> Result, DataFusionError> { + use iceberg_rust::spec::partition::Transform; + // Skip identity-transform partitions whose partition name matches the + // source column name. For Iceberg identity partitions on existing columns + // (e.g. `identity(event_name)`), the column is present in both the parquet + // file body and the Hive-style directory path. datafusion's parquet reader + // subtracts matching partition cols from the file schema when computing + // the expected column count, which causes an off-by-one mismatch + // ("expected N cols but got N+1"). Transformed partitions (day, hour, etc.) + // and renamed-identity partitions keep their own synthetic column. let table_partition_cols: Vec = partition_fields .iter() + .filter(|partition_field| { + !(matches!(partition_field.transform(), Transform::Identity) + && partition_field.name() == partition_field.source_name()) + }) .map(|partition_field| { Ok(Field::new( partition_field.name().to_owned(), diff --git a/iceberg-rust-spec/src/spec/manifest_list.rs b/iceberg-rust-spec/src/spec/manifest_list.rs index 4d47f8b0..727d479f 100644 --- a/iceberg-rust-spec/src/spec/manifest_list.rs +++ b/iceberg-rust-spec/src/spec/manifest_list.rs @@ -130,10 +130,13 @@ mod _serde { /// ID of the snapshot where the manifest file was added pub added_snapshot_id: i64, /// Number of entries in the manifest that have status ADDED (1), when null this is assumed to be non-zero + #[serde(rename = "added_data_files_count", alias = "added_files_count")] pub added_files_count: i32, /// Number of entries in the manifest that have status EXISTING (0), when null this is assumed to be non-zero + #[serde(rename = "existing_data_files_count", alias = "existing_files_count")] pub existing_files_count: i32, /// Number of entries in the manifest that have status DELETED (2), when null this is assumed to be non-zero + #[serde(rename = "deleted_data_files_count", alias = "deleted_files_count")] pub deleted_files_count: i32, /// Number of rows in all of files in the manifest that have status ADDED, when null this is assumed to be non-zero pub added_rows_count: i64, @@ -570,17 +573,17 @@ pub fn manifest_list_schema_v2() -> &'static AvroSchema { "field-id": 503 }, { - "name": "added_files_count", + "name": "added_data_files_count", "type": "int", "field-id": 504 }, { - "name": "existing_files_count", + "name": "existing_data_files_count", "type": "int", "field-id": 505 }, { - "name": "deleted_files_count", + "name": "deleted_data_files_count", "type": "int", "field-id": 506 }, @@ -843,4 +846,101 @@ mod tests { ); } } + + /// Simulates reading a v2 manifest list written by Apache Iceberg >= 1.0, + /// which uses the renamed field names `added_data_files_count`, + /// `existing_data_files_count`, `deleted_data_files_count`. Our reader + /// schema declares these as Avro field aliases, so resolution should succeed. + #[test] + pub fn test_manifest_list_v2_apache_field_names() { + // Writer schema using Apache Iceberg spec-v2 field names. + let apache_writer_schema = apache_avro::Schema::parse_str( + r#" + { + "type": "record", + "name": "manifest_file", + "fields": [ + {"name": "manifest_path", "type": "string", "field-id": 500}, + {"name": "manifest_length", "type": "long", "field-id": 501}, + {"name": "partition_spec_id", "type": "int", "field-id": 502}, + {"name": "content", "type": "int", "field-id": 517}, + {"name": "sequence_number", "type": "long", "field-id": 515}, + {"name": "min_sequence_number", "type": "long", "field-id": 516}, + {"name": "added_snapshot_id", "type": "long", "field-id": 503}, + {"name": "added_data_files_count", "type": "int", "field-id": 504}, + {"name": "existing_data_files_count", "type": "int", "field-id": 505}, + {"name": "deleted_data_files_count", "type": "int", "field-id": 506}, + {"name": "added_rows_count", "type": "long", "field-id": 512}, + {"name": "existing_rows_count", "type": "long", "field-id": 513}, + {"name": "deleted_rows_count", "type": "long", "field-id": 514}, + { + "name": "partitions", + "type": ["null", { + "type": "array", + "items": { + "type": "record", + "name": "r508", + "fields": [ + {"name": "contains_null", "type": "boolean", "field-id": 509}, + {"name": "contains_nan", "type": ["null", "boolean"], "field-id": 518}, + {"name": "lower_bound", "type": ["null", "bytes"], "field-id": 510}, + {"name": "upper_bound", "type": ["null", "bytes"], "field-id": 511} + ] + }, + "element-id": 508 + }], + "default": null, + "field-id": 507 + }, + { + "name": "key_metadata", + "type": ["null", "bytes"], + "default": null, + "field-id": 519 + } + ] + } + "#, + ) + .unwrap(); + + // Build an Avro record using the Apache field names. + let mut record = apache_avro::types::Record::new(&apache_writer_schema).unwrap(); + record.put("manifest_path", "s3://bucket/path/manifest-0.avro"); + record.put("manifest_length", 1200_i64); + record.put("partition_spec_id", 0_i32); + record.put("content", 0_i32); + record.put("sequence_number", 566_i64); + record.put("min_sequence_number", 0_i64); + record.put("added_snapshot_id", 39487483032_i64); + record.put("added_data_files_count", 1_i32); + record.put("existing_data_files_count", 2_i32); + record.put("deleted_data_files_count", 0_i32); + record.put("added_rows_count", 1000_i64); + record.put("existing_rows_count", 8000_i64); + record.put("deleted_rows_count", 0_i64); + record.put("partitions", AvroValue::Union(0, Box::new(AvroValue::Null))); + record.put("key_metadata", AvroValue::Union(0, Box::new(AvroValue::Null))); + + let mut writer = apache_avro::Writer::new(&apache_writer_schema, Vec::new()); + writer.append(record).unwrap(); + let encoded = writer.into_inner().unwrap(); + + // Read with iceberg-rust's schema (aliases declared). + let reader_schema = manifest_list_schema_v2(); + let reader = apache_avro::Reader::with_schema(reader_schema, &*encoded).unwrap(); + + let mut count = 0; + for record in reader { + let result = + apache_avro::from_value::<_serde::ManifestListEntryV2>(&record.unwrap()).unwrap(); + assert_eq!(result.added_files_count, 1); + assert_eq!(result.existing_files_count, 2); + assert_eq!(result.deleted_files_count, 0); + assert_eq!(result.added_rows_count, 1000); + assert_eq!(result.manifest_path, "s3://bucket/path/manifest-0.avro"); + count += 1; + } + assert_eq!(count, 1, "expected exactly one record"); + } } From b9404596918e024be5cf5384dac900aab0e56f59 Mon Sep 17 00:00:00 2001 From: Sergei Turukin Date: Wed, 15 Apr 2026 15:13:41 -0500 Subject: [PATCH 2/6] refactor(table_scan): promote identity-self-named partition filter to table_scan Move the identity-self-named partition drop out of `datafusion_partition_columns` and into `table_scan` itself, so the physical scan column set and the manifest pruner's `partition_column_names` set are computed from the same filtered list. Previously they diverged: `datafusion_partition_columns` filtered out identity-self-named fields but the pruner still built its column subset from the unfiltered `partition_fields`, which meant filters on identity-self-named columns were incorrectly routed through `PruneManifests` (and then failed the subset test on the reduced partition schema anyway). Introduces `file_partition_fields` (kept) and `drop_partition_indices` (dropped), constructed once at the top of `table_scan` from the unfiltered `partition_fields`. Both are then threaded through every downstream consumer: - `datafusion_partition_columns` is called with the kept list. - The manifest-level pruner's `partition_column_names` set is built from the kept list and a comment documents that identity-self-named predicates are intentionally excluded here because they are pruned by per-file statistics in `PruneDataFiles` instead. - `drop_partition_indices` is later consumed by `generate_partitioned_file` so callers that still need to see the unfiltered partition-field order can account for the gaps. Prerequisite for follow-up commits that add the projection remap, TIMESTAMPTZ transform acceptance, PruneDataFiles arrow_schema fix, and manifest nested-id resolution. --- datafusion_iceberg/src/table.rs | 66 +++++++++++++++++++++++---------- 1 file changed, 46 insertions(+), 20 deletions(-) diff --git a/datafusion_iceberg/src/table.rs b/datafusion_iceberg/src/table.rs index 67d989bf..7358664a 100644 --- a/datafusion_iceberg/src/table.rs +++ b/datafusion_iceberg/src/table.rs @@ -432,7 +432,29 @@ async fn table_scan( None }; - let mut table_partition_cols = datafusion_partition_columns(partition_fields)?; + // Compute the subset of partition fields that are materialized via the + // Hive-style directory path (i.e. NOT already present in the parquet file + // body). Iceberg identity partitions on source columns duplicate those + // columns into both the file and the path; for the datafusion scan we must + // omit them from `table_partition_cols` so the parquet reader doesn't try + // to project them from the path. Column-level pruning still works for + // identity partitions via the per-file lower/upper bounds in + // PruneDataFiles, so we lose nothing by excluding them here. + let (file_partition_fields, drop_partition_indices): (Vec<&BoundPartitionField<'_>>, Vec) = { + use iceberg_rust::spec::partition::Transform; + let mut kept: Vec<&BoundPartitionField<'_>> = Vec::new(); + let mut dropped: Vec = Vec::new(); + for (i, pf) in partition_fields.iter().enumerate() { + if matches!(pf.transform(), Transform::Identity) && pf.name() == pf.source_name() { + dropped.push(i); + } else { + kept.push(pf); + } + } + (kept, dropped) + }; + + let mut table_partition_cols = datafusion_partition_columns(&file_partition_fields)?; let file_schema: SchemaRef = Arc::new((schema.fields()).try_into().unwrap()); @@ -472,11 +494,14 @@ async fn table_scan( physical_predicate.clone() { let partition_schema = Arc::new(ArrowSchema::new(table_partition_cols.clone())); - let partition_column_names = partition_fields + // Use only the file-path-materialized partition fields here so that + // predicates on identity-self-named columns (which are pruned via + // PruneDataFiles) are NOT routed through the manifest-level + // partition pruner. + let partition_column_names = file_partition_fields .iter() - .map(|field| Ok(field.source_name().to_owned())) - .collect::, Error>>() - .map_err(DataFusionIcebergError::from)?; + .map(|field| field.source_name().to_owned()) + .collect::>(); let partition_predicates = conjunction( filters @@ -633,6 +658,7 @@ async fn table_scan( .then(|(partition_value, mut delete_files)| { let object_store_url = object_store_url.clone(); let table_partition_cols = table_partition_cols.clone(); + let drop_partition_indices = drop_partition_indices.clone(); let statistics = statistics.clone(); let physical_predicate = physical_predicate.clone(); let schema = &schema; @@ -689,6 +715,7 @@ async fn table_scan( .try_fold(None, |acc, delete_manifest| { let object_store_url = object_store_url.clone(); let table_partition_cols = table_partition_cols.clone(); + let drop_partition_indices = drop_partition_indices.clone(); let statistics = statistics.clone(); let physical_predicate = physical_predicate.clone(); let schema = &schema; @@ -713,6 +740,7 @@ async fn table_scan( last_updated_ms, enable_data_file_path_column, manifest_path, + &drop_partition_indices, ) .unwrap(); data_files.push(data_file); @@ -741,6 +769,7 @@ async fn table_scan( last_updated_ms, enable_data_file_path_column, manifest_path, + &drop_partition_indices, )?; let delete_file_source = Arc::new( @@ -843,6 +872,7 @@ async fn table_scan( last_updated_ms, enable_data_file_path_column, manifest_path, + &drop_partition_indices, ) }) .collect::, _>>()?; @@ -892,6 +922,7 @@ async fn table_scan( last_updated_ms, enable_data_file_path_column, manifest_path, + &drop_partition_indices, ) .unwrap() }) @@ -930,23 +961,10 @@ async fn table_scan( } fn datafusion_partition_columns( - partition_fields: &[BoundPartitionField<'_>], + partition_fields: &[&BoundPartitionField<'_>], ) -> Result, DataFusionError> { - use iceberg_rust::spec::partition::Transform; - // Skip identity-transform partitions whose partition name matches the - // source column name. For Iceberg identity partitions on existing columns - // (e.g. `identity(event_name)`), the column is present in both the parquet - // file body and the Hive-style directory path. datafusion's parquet reader - // subtracts matching partition cols from the file schema when computing - // the expected column count, which causes an off-by-one mismatch - // ("expected N cols but got N+1"). Transformed partitions (day, hour, etc.) - // and renamed-identity partitions keep their own synthetic column. let table_partition_cols: Vec = partition_fields .iter() - .filter(|partition_field| { - !(matches!(partition_field.transform(), Transform::Identity) - && partition_field.name() == partition_field.source_name()) - }) .map(|partition_field| { Ok(Field::new( partition_field.name().to_owned(), @@ -1051,13 +1069,21 @@ fn generate_partitioned_file( last_updated_ms: i64, enable_data_file_path: bool, manifest_file_path: Option, + drop_partition_indices: &[usize], ) -> Result { let manifest_statistics = manifest_statistics(schema, manifest); + // Iceberg stores partition values in the order of the partition spec + // (one entry per partition field). Drop entries that correspond to + // identity-self-named partitions we've excluded from `table_partition_cols` + // so that `file.partition_values.len()` matches the filtered + // `partition_fields` length passed to datafusion's `FilePruner`. let mut partition_values = manifest .data_file() .partition() .iter() - .map(|x| { + .enumerate() + .filter(|(i, _)| !drop_partition_indices.contains(i)) + .map(|(_, x)| { x.as_ref() .map(value_to_scalarvalue) .unwrap_or(Ok(ScalarValue::Null)) From 690b42b4d94c609699a0ad95156b660dda840506 Mon Sep 17 00:00:00 2001 From: Sergei Turukin Date: Tue, 14 Apr 2026 20:00:59 -0500 Subject: [PATCH 3/6] fix(datafusion_iceberg): remap caller projection to combined-schema space DataFusionTable::schema() returns [file_schema, __data_file_path?, __manifest_file_path?] but the physical FileScanConfig output is [file_schema, kept_partition_transform_cols..., __data_file_path?, __manifest_file_path?]. Any partition spec with a non-identity transform (day, hour, month, year, bucket, truncate) creates synthetic columns (e.g. ts_day for day(ts)) that sit between the user columns and the metadata columns in the physical schema but are absent from the provider schema. table_scan() was passing the caller's `projection` (indices into the provider schema) straight through to FileScanConfig::with_projection, which interprets indices against the combined schema. With enable_data_file_path_column=true this picked up `ts_day` at the slot where `__data_file_path` was expected and silently truncated `__manifest_file_path`, which in turn made any downstream ProjectionExec referencing those columns by name+index fail with: Internal error: Input field name _ does not match with the projection expression __data_file_path. Embucket's MERGE COW planner hits this on every partitioned target. Fix: compute combined_projection once from the caller's projection, remapping provider-schema indices for __data_file_path / __manifest_file_path to their actual positions in [file_schema, kept_partition_cols, __data_file_path?, __manifest_file_path?]. Use combined_projection throughout table_scan (no-delete path, equality-delete base, per-closure clones). Adds 7 regression tests (day, hour, month, year, bucket, truncate, renamed-identity) alongside the existing unpartitioned test_datafusion_table_insert_with_data_file_path. Co-Authored-By: Claude Opus 4.6 (1M context) --- datafusion_iceberg/src/table.rs | 433 +++++++++++++++++++++++++++++++- 1 file changed, 421 insertions(+), 12 deletions(-) diff --git a/datafusion_iceberg/src/table.rs b/datafusion_iceberg/src/table.rs index 7358664a..d79edb5b 100644 --- a/datafusion_iceberg/src/table.rs +++ b/datafusion_iceberg/src/table.rs @@ -458,11 +458,62 @@ async fn table_scan( let file_schema: SchemaRef = Arc::new((schema.fields()).try_into().unwrap()); - // If no projection was specified default to projecting all the fields + if enable_data_file_path_column { + table_partition_cols.push(Field::new(DATA_FILE_PATH_COLUMN, DataType::Utf8, false)); + } + + if enable_manifest_file_path_column { + table_partition_cols.push(Field::new(MANIFEST_FILE_PATH_COLUMN, DataType::Utf8, false)); + } + + // If no projection was specified default to projecting all the fields of + // the provider-facing (arrow) schema, which already contains the optional + // `__data_file_path` / `__manifest_file_path` columns. let projection = projection .cloned() .unwrap_or((0..arrow_schema.fields().len()).collect_vec()); + // DataFusion's `FileScanConfig.with_projection` interprets indices against + // the combined schema `[file_schema, table_partition_cols]`. Our caller, + // however, projects against `arrow_schema` (the `DataFusionTable::schema`) + // which is `[file_schema, __data_file_path?, __manifest_file_path?]` — + // synthetic partition-transform columns like `ts_day` for `day(ts)` are + // *not* in `arrow_schema`. When those synthetic columns exist, the two + // schemas diverge and passing the caller's indices through unmodified + // picks up `ts_day` where `__data_file_path` was expected, shifting every + // metadata column and silently truncating `__manifest_file_path`. It + // triggers DataFusion's "Input field name _ does + // not match with the projection expression __data_file_path" error as + // soon as a downstream ProjectionExec tries to reference those columns + // by name (e.g. from Embucket's MERGE COW planner). + // + // Fix: remap the caller's projection from provider-schema space to + // combined-schema space. + let file_schema_len = file_schema.fields().len(); + let kept_partition_len = file_partition_fields.len(); + let combined_projection: Vec = projection + .iter() + .map(|&idx| { + if idx < file_schema_len { + idx + } else { + let name = arrow_schema.fields[idx].name().as_str(); + if name == DATA_FILE_PATH_COLUMN && enable_data_file_path_column { + file_schema_len + kept_partition_len + } else if name == MANIFEST_FILE_PATH_COLUMN && enable_manifest_file_path_column { + file_schema_len + + kept_partition_len + + usize::from(enable_data_file_path_column) + } else { + // Fallback: preserve the original index. This only fires + // for unexpected columns and will surface as a clear + // error further down the planning pipeline. + idx + } + } + }) + .collect(); + let projection_expr: Vec<_> = projection .iter() .enumerate() @@ -475,14 +526,6 @@ async fn table_scan( }) .collect(); - if enable_data_file_path_column { - table_partition_cols.push(Field::new(DATA_FILE_PATH_COLUMN, DataType::Utf8, false)); - } - - if enable_manifest_file_path_column { - table_partition_cols.push(Field::new(MANIFEST_FILE_PATH_COLUMN, DataType::Utf8, false)); - } - // All files have to be grouped according to their partition values. This is done by using a HashMap with the partition values as the key. // This way data files with the same partition value are mapped to the same vector. let mut data_file_groups: HashMap> = HashMap::new(); @@ -665,7 +708,7 @@ async fn table_scan( let file_schema = file_schema.clone(); let file_source = file_source.clone(); let projection_expr = projection_expr.clone(); - let projection = &projection; + let combined_projection = combined_projection.clone(); let mut data_files = data_file_groups .remove(&partition_value) .unwrap_or_default(); @@ -690,7 +733,10 @@ async fn table_scan( // each equality delete file may have different deletion columns. // And since we need to reconcile them all with data files using joins and unions, // we need to make sure their schemas are fully compatible in all intermediate nodes. - let mut equality_projection = projection.clone(); + // Indices are in combined-schema space (see `combined_projection` above); the + // delete-id lookup below adds user-column indices, which already match file_schema + // positions 1:1 so no remap is needed for the appended entries. + let mut equality_projection = combined_projection.clone(); delete_files .iter() .flat_map(|delete_manifest| delete_manifest.1.data_file().equality_ids()) @@ -935,7 +981,7 @@ async fn table_scan( FileScanConfigBuilder::new(object_store_url, file_schema, file_source) .with_file_groups(file_groups) .with_statistics(statistics) - .with_projection(Some(projection.clone())) + .with_projection(Some(combined_projection.clone())) .with_limit(limit) .with_table_partition_cols(table_partition_cols) .build(); @@ -2594,6 +2640,369 @@ mod tests { } } + /// Regression test for the MERGE-on-partitioned-table bug. + /// + /// When a table has a non-identity partition transform (day, hour, month, + /// year, bucket, truncate), `datafusion_partition_columns()` pushes a + /// synthetic partition-transform column (e.g. `ts_day` for `day(ts)`) into + /// `table_partition_cols`. That column ends up in the physical scan output + /// between the user columns and the `__data_file_path` / `__manifest_file_path` + /// columns. But `DataFusionTable::schema()` (the logical schema) only + /// contains user columns + `__data_file_path` + `__manifest_file_path` — + /// no synthetic partition column — so DataFusion validates the projection + /// against a physical output that has the wrong field at the position where + /// `__data_file_path` is expected, yielding: + /// + /// Internal error: Input field name ts_day does not match with the + /// projection expression __data_file_path. + /// + /// This test exercises the buggy path with a `day(ts)` partition spec. It + /// should pass: the scan must emit exactly the DataFusionTable::schema() + /// columns in the expected order. + #[tokio::test] + pub async fn test_data_file_path_with_day_partitioned_table() { + let object_store = ObjectStoreBuilder::memory(); + + let catalog: Arc = Arc::new( + SqlCatalog::new("sqlite://", "test", object_store) + .await + .unwrap(), + ); + + let schema = Schema::builder() + .with_struct_field(StructField { + id: 1, + name: "id".to_string(), + required: true, + field_type: Type::Primitive(PrimitiveType::Long), + doc: None, + }) + .with_struct_field(StructField { + id: 2, + name: "ts".to_string(), + required: true, + field_type: Type::Primitive(PrimitiveType::Date), + doc: None, + }) + .with_struct_field(StructField { + id: 3, + name: "payload".to_string(), + required: false, + field_type: Type::Primitive(PrimitiveType::String), + doc: None, + }) + .build() + .unwrap(); + + let partition_spec = PartitionSpec::builder() + .with_partition_field(PartitionField::new(2, 1000, "ts_day", Transform::Day)) + .build() + .expect("Failed to build partition spec"); + + let table = Table::builder() + .with_name("events") + .with_location("/test/events") + .with_schema(schema) + .with_partition_spec(partition_spec) + .build(&["test".to_owned()], catalog) + .await + .expect("Failed to create partitioned table"); + + let config = crate::table::DataFusionTableConfigBuilder::default() + .enable_data_file_path_column(true) + .enable_manifest_file_path_column(true) + .build() + .unwrap(); + + let table = Arc::new(DataFusionTable::new_with_config( + Tabular::Table(table), + None, + None, + None, + Some(config), + )); + + let ctx = SessionContext::new(); + ctx.register_table("events", table.clone()).unwrap(); + + ctx.sql( + "INSERT INTO events (id, ts, payload) VALUES + (1, DATE '2020-01-01', 'a'), + (2, DATE '2020-01-02', 'b'), + (3, DATE '2020-01-03', 'c');", + ) + .await + .expect("Failed to create query plan for insert") + .collect() + .await + .expect("Failed to insert values into partitioned table"); + + let batches = ctx + .sql("select * from events;") + .await + .expect("Failed to create plan for select") + .collect() + .await + .expect("Failed to execute select query"); + + // The scan's output schema must match DataFusionTable::schema(). In + // particular: no synthetic `ts_day` column leaking through, and + // `__data_file_path` sitting where the logical schema expects it. + let provider_schema: Vec = table + .schema + .fields() + .iter() + .map(|f: &Arc| f.name().to_owned()) + .collect(); + assert_eq!( + provider_schema, + vec![ + "id".to_string(), + "ts".to_string(), + "payload".to_string(), + "__data_file_path".to_string(), + "__manifest_file_path".to_string(), + ], + "DataFusionTable::schema() precondition" + ); + + let mut total_rows = 0; + let mut saw_non_empty = false; + for batch in &batches { + if batch.num_rows() == 0 { + continue; + } + saw_non_empty = true; + total_rows += batch.num_rows(); + + let actual_names: Vec = batch + .schema() + .fields() + .iter() + .map(|f| f.name().to_owned()) + .collect(); + assert_eq!( + actual_names, provider_schema, + "scan output schema must match provider schema; \ + synthetic partition columns (e.g. ts_day) must not leak through" + ); + + let data_file_path_column = batch + .column_by_name("__data_file_path") + .expect("__data_file_path column should exist in scan output"); + let values = data_file_path_column + .as_any() + .downcast_ref::() + .expect("__data_file_path should be a StringArray"); + for i in 0..batch.num_rows() { + let value = values.value(i); + assert!( + !value.is_empty(), + "__data_file_path should not be empty for row {i}" + ); + assert!( + value.contains(".parquet"), + "__data_file_path should contain .parquet, got {value}" + ); + } + } + + assert!(saw_non_empty, "expected at least one non-empty batch"); + assert_eq!(total_rows, 3, "expected 3 rows total"); + } + + /// Parameterised harness exercising every partition transform against + /// the `enable_data_file_path_column` scan path. Each transform should + /// yield the same provider schema (`id`, `ts`, `kind`, `__data_file_path`, + /// `__manifest_file_path`) regardless of whether it rewrites the + /// partition field name. + async fn run_data_file_path_transform_case( + case_name: &'static str, + partition_source_id: i32, + partition_name: &'static str, + transform: Transform, + ) { + let object_store = ObjectStoreBuilder::memory(); + + let catalog: Arc = Arc::new( + SqlCatalog::new("sqlite://", case_name, object_store) + .await + .unwrap(), + ); + + let schema = Schema::builder() + .with_struct_field(StructField { + id: 1, + name: "id".to_string(), + required: true, + field_type: Type::Primitive(PrimitiveType::Long), + doc: None, + }) + .with_struct_field(StructField { + id: 2, + name: "ts".to_string(), + required: true, + field_type: Type::Primitive(PrimitiveType::Timestamp), + doc: None, + }) + .with_struct_field(StructField { + id: 3, + name: "kind".to_string(), + required: true, + field_type: Type::Primitive(PrimitiveType::String), + doc: None, + }) + .build() + .unwrap(); + + let partition_spec = PartitionSpec::builder() + .with_partition_field(PartitionField::new( + partition_source_id, + 1000, + partition_name, + transform, + )) + .build() + .unwrap_or_else(|e| panic!("[{case_name}] failed to build spec: {e}")); + + let table = Table::builder() + .with_name(case_name) + .with_location(&format!("/test/{case_name}")) + .with_schema(schema) + .with_partition_spec(partition_spec) + .build(&["test".to_owned()], catalog) + .await + .unwrap_or_else(|e| panic!("[{case_name}] failed to create table: {e}")); + + let config = crate::table::DataFusionTableConfigBuilder::default() + .enable_data_file_path_column(true) + .enable_manifest_file_path_column(true) + .build() + .unwrap(); + + let table = Arc::new(DataFusionTable::new_with_config( + Tabular::Table(table), + None, + None, + None, + Some(config), + )); + + let ctx = SessionContext::new(); + ctx.register_table(case_name, table.clone()).unwrap(); + + ctx.sql(&format!( + "INSERT INTO {case_name} (id, ts, kind) VALUES + (1, TIMESTAMP '2020-01-01 00:00:00', 'a'), + (2, TIMESTAMP '2020-01-02 12:00:00', 'b'), + (3, TIMESTAMP '2020-02-03 06:30:00', 'c');" + )) + .await + .unwrap_or_else(|e| panic!("[{case_name}] insert plan failed: {e}")) + .collect() + .await + .unwrap_or_else(|e| panic!("[{case_name}] insert exec failed: {e}")); + + let batches = ctx + .sql(&format!("select * from {case_name};")) + .await + .unwrap_or_else(|e| panic!("[{case_name}] select plan failed: {e}")) + .collect() + .await + .unwrap_or_else(|e| panic!("[{case_name}] select exec failed: {e}")); + + let expected: Vec = table + .schema + .fields() + .iter() + .map(|f: &Arc| f.name().to_owned()) + .collect(); + assert_eq!( + expected, + vec![ + "id".to_string(), + "ts".to_string(), + "kind".to_string(), + "__data_file_path".to_string(), + "__manifest_file_path".to_string(), + ], + "[{case_name}] provider schema precondition" + ); + + let mut total_rows = 0; + for batch in &batches { + if batch.num_rows() == 0 { + continue; + } + total_rows += batch.num_rows(); + let actual: Vec = batch + .schema() + .fields() + .iter() + .map(|f: &Arc| f.name().to_owned()) + .collect(); + assert_eq!( + actual, expected, + "[{case_name}] scan output schema must match provider schema" + ); + assert!( + batch + .column_by_name("__data_file_path") + .expect("__data_file_path missing") + .as_any() + .downcast_ref::() + .expect("__data_file_path should be StringArray") + .iter() + .all(|v| v.is_some_and(|s| s.contains(".parquet"))), + "[{case_name}] __data_file_path values must be populated" + ); + } + assert_eq!(total_rows, 3, "[{case_name}] expected 3 rows"); + } + + #[tokio::test] + pub async fn test_data_file_path_with_hour_partition() { + run_data_file_path_transform_case("hour_probe", 2, "ts_hour", Transform::Hour).await; + } + + #[tokio::test] + pub async fn test_data_file_path_with_month_partition() { + run_data_file_path_transform_case("month_probe", 2, "ts_month", Transform::Month).await; + } + + #[tokio::test] + pub async fn test_data_file_path_with_year_partition() { + run_data_file_path_transform_case("year_probe", 2, "ts_year", Transform::Year).await; + } + + #[tokio::test] + pub async fn test_data_file_path_with_bucket_partition() { + run_data_file_path_transform_case("bucket_probe", 1, "id_bucket", Transform::Bucket(4)) + .await; + } + + #[tokio::test] + pub async fn test_data_file_path_with_truncate_partition() { + // Truncate on the `id` Long column. Truncate on strings isn't + // supported by the current iceberg-rust transform implementation. + run_data_file_path_transform_case("trunc_probe", 1, "id_trunc", Transform::Truncate(10)) + .await; + } + + #[tokio::test] + pub async fn test_data_file_path_with_identity_partition_renamed() { + // Identity with a renamed partition field (different name from source column) + // also goes through `datafusion_partition_columns()` because the + // identity-self-named drop only fires when pf.name() == pf.source_name(). + run_data_file_path_transform_case( + "identity_renamed_probe", + 3, + "kind_id", + Transform::Identity, + ) + .await; + } + #[test] fn test_fake_object_store_url() { assert_eq!( From 8d242d56c25b58e209362cff5307358ed1944559 Mon Sep 17 00:00:00 2001 From: Sergei Turukin Date: Tue, 14 Apr 2026 20:45:48 -0500 Subject: [PATCH 4/6] fix(arrow/transform): accept TIMESTAMP_TZ for day/hour/month/year transforms MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `transform_arrow()` only matched `DataType::Timestamp(TimeUnit::Microsecond, None)` for the day/hour/month/year arms, so any `timestamptz` column fell through to the catchall and raised `Compute error: Failed to perform transform for datatype`. Embucket's MERGE write path on `events_hooli` — whose `collector_tstamp` is `TIMESTAMP_TZ` partitioned by `day(collector_tstamp)` — tripped this every time. Iceberg's day/hour/month/year transforms are defined on the absolute instant (microseconds since the Unix epoch), so the Arrow timezone metadata is irrelevant to the numeric result. Widen each arm to `Timestamp(Microsecond, _)`. For month and year the existing `date_part` call used a named-timezone path that requires `chrono-tz`; cast to `Timestamp(Microsecond, None)` first so we run on a naive variant that works without that feature flag. Adds 4 regression tests exercising all four transforms with a `TimestampMicrosecondArray::with_timezone("UTC")` input to lock the fix in. Co-Authored-By: Claude Opus 4.6 (1M context) --- iceberg-rust/src/arrow/transform.rs | 97 ++++++++++++++++++++++++++--- 1 file changed, 90 insertions(+), 7 deletions(-) diff --git a/iceberg-rust/src/arrow/transform.rs b/iceberg-rust/src/arrow/transform.rs index b7f93675..1305a43a 100644 --- a/iceberg-rust/src/arrow/transform.rs +++ b/iceberg-rust/src/arrow/transform.rs @@ -72,25 +72,36 @@ pub fn transform_arrow(array: ArrayRef, transform: &Transform) -> Result { + // `_` for the timezone parameter so both `timestamp` (None) and + // `timestamptz` (Some(tz)) match. Iceberg partition transforms are + // defined on the absolute instant (microseconds since Unix epoch in + // UTC), so the attached tz metadata is irrelevant to the numeric + // result — we just need to read the underlying i64. + (DataType::Timestamp(TimeUnit::Microsecond, _), Transform::Hour) => { Ok(Arc::new(unary::<_, _, Int32Type>( as_primitive_array::(&cast(&array, &DataType::Int64)?), micros_to_hours, )) as Arc) } - (DataType::Timestamp(TimeUnit::Microsecond, None), Transform::Day) => { + (DataType::Timestamp(TimeUnit::Microsecond, _), Transform::Day) => { Ok(Arc::new(unary::<_, _, Int32Type>( as_primitive_array::(&cast(&array, &DataType::Int64)?), micros_to_days, )) as Arc) } - (DataType::Timestamp(TimeUnit::Microsecond, None), Transform::Month) => { + (DataType::Timestamp(TimeUnit::Microsecond, _), Transform::Month) => { + // date_part requires chrono-tz for named timezones like "UTC". + // Iceberg computes month/year over the absolute instant, so the + // tz metadata only affects display, not the result. Strip the tz + // by casting to Timestamp(Microsecond, None) first so date_part + // runs on a plain value. + let naive = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None))?; let year = date_part( - as_primitive_array::(&array), + as_primitive_array::(&naive), DatePart::Year, )?; let month = date_part( - as_primitive_array::(&array), + as_primitive_array::(&naive), DatePart::Month, )?; Ok(Arc::new(binary::<_, _, _, Int32Type>( @@ -99,10 +110,12 @@ pub fn transform_arrow(array: ArrayRef, transform: &Transform) -> Result { + (DataType::Timestamp(TimeUnit::Microsecond, _), Transform::Year) => { + // Same tz-stripping rationale as Month above. + let naive = cast(&array, &DataType::Timestamp(TimeUnit::Microsecond, None))?; Ok(Arc::new(unary::<_, _, Int32Type>( as_primitive_array::(&date_part( - as_primitive_array::(&array), + as_primitive_array::(&naive), DatePart::Year, )?), datepart_to_years, @@ -520,4 +533,74 @@ mod tests { "Compute error: Failed to perform transform for datatype" ); } + + /// Returns the same three representative microsecond values as + /// `create_timestamp_micro_array()` but wrapped in a `TimestampMicrosecondArray` + /// with a `"UTC"` timezone attached. This matches how Arrow encodes + /// Iceberg's `timestamptz` type, which is what Embucket hands to the + /// partition transforms for tables like `events_hooli` whose + /// `collector_tstamp` is `TIMESTAMP_TZ`. + fn create_timestamp_micro_tz_array() -> ArrayRef { + Arc::new( + TimestampMicrosecondArray::from(vec![ + Some(1682937000000000), + Some(1686840330000000), + Some(1704067200000000), + None, + ]) + .with_timezone("UTC"), + ) as ArrayRef + } + + #[test] + fn test_timestamp_tz_day_transform() { + let array = create_timestamp_micro_tz_array(); + let result = transform_arrow(array, &Transform::Day).unwrap(); + let expected = Arc::new(arrow::array::Int32Array::from(vec![ + Some(19478), + Some(19523), + Some(19723), + None, + ])) as ArrayRef; + assert_eq!(&expected, &result); + } + + #[test] + fn test_timestamp_tz_hour_transform() { + let array = create_timestamp_micro_tz_array(); + let result = transform_arrow(array, &Transform::Hour).unwrap(); + let expected = Arc::new(arrow::array::Int32Array::from(vec![ + Some(467482), + Some(468566), + Some(473352), + None, + ])) as ArrayRef; + assert_eq!(&expected, &result); + } + + #[test] + fn test_timestamp_tz_month_transform() { + let array = create_timestamp_micro_tz_array(); + let result = transform_arrow(array, &Transform::Month).unwrap(); + let expected = Arc::new(arrow::array::Int32Array::from(vec![ + Some(641), + Some(642), + Some(649), + None, + ])) as ArrayRef; + assert_eq!(&expected, &result); + } + + #[test] + fn test_timestamp_tz_year_transform() { + let array = create_timestamp_micro_tz_array(); + let result = transform_arrow(array, &Transform::Year).unwrap(); + let expected = Arc::new(arrow::array::Int32Array::from(vec![ + Some(53), + Some(53), + Some(54), + None, + ])) as ArrayRef; + assert_eq!(&expected, &result); + } } From cac01cd4a7e872a0b690d049fddfefa045449a8b Mon Sep 17 00:00:00 2001 From: Sergei Turukin Date: Wed, 15 Apr 2026 13:53:49 -0500 Subject: [PATCH 5/6] fix(datafusion_iceberg): route full arrow_schema to PruneDataFiles MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The second-stage data-file pruner (PruneDataFiles) was constructed with `partition_schema` — a subset schema holding only the Hive-style partition columns. Its `min_values`/`max_values` implementation looks up each column referenced by the pruning predicate via `arrow_schema.field_with_name(..)` to fetch the datatype, so any filter on a column absent from `partition_schema` silently returned `None` and pruned nothing. Identity-self-named partition columns (where `pf.name() == pf.source_name()`) are intentionally dropped from `file_partition_fields` so the parquet reader doesn't duplicate them between the path encoding and the file body, which also drops them from `table_partition_cols` and therefore from `partition_schema`. The result: a filter like `event_name = 'ad_start'` against a table partitioned by `identity(event_name)` reached the second- stage pruner but found no schema hit, so every partition file of the target was scanned in full (`files_ranges_pruned_statistics=0`). This only surfaced now because Embucket/embucket#126 unblocked the filter reaching TableScan in the first place. Fix: pass the full `arrow_schema` to `PruneDataFiles::new`. It has every column the predicate might reference — identity-self-named partition columns, non-partition columns with per-file statistics, etc. Correctness is preserved because the first-stage `PruneManifests` path still prunes transformed partition columns (`collector_tstamp_day`, `id_bucket`, ...) via manifest-list partition bounds, and synthetic partition-transform columns simply return `None` from `PruneDataFiles` (no per-file stats exist for them), which is the same behavior they had before. Adds a regression test: `test_identity_self_named_partition_filter_prunes_files` creates a `identity(kind)` partitioned table, inserts one row per partition value to materialize 3 distinct parquet files, then scans with `kind = 'a'` and asserts the resulting plan lists exactly 1 parquet file instead of 3. Refs: Embucket/embucket#127 Co-Authored-By: Claude Opus 4.6 (1M context) --- datafusion_iceberg/src/table.rs | 124 +++++++++++++++++++++++++++++++- 1 file changed, 122 insertions(+), 2 deletions(-) diff --git a/datafusion_iceberg/src/table.rs b/datafusion_iceberg/src/table.rs index d79edb5b..4ecda629 100644 --- a/datafusion_iceberg/src/table.rs +++ b/datafusion_iceberg/src/table.rs @@ -597,10 +597,19 @@ async fn table_scan( let pruning_predicate = PruningPredicate::try_new(physical_predicate, arrow_schema.clone())?; - // After the first pruning stage the data_files are pruned again based on the pruning statistics in the manifest files. + // After the first pruning stage the data_files are pruned again based + // on the pruning statistics in the manifest files. `PruneDataFiles` + // looks up each column referenced by the predicate in its + // `arrow_schema` field to fetch the datatype; passing the narrow + // `partition_schema` here would hide every non-partition-key column + // (including identity-self-named partition columns that have been + // dropped from `table_partition_cols` because they are materialized + // in the parquet file body), so any filter on such a column would + // silently prune nothing. Passing the full `arrow_schema` lets the + // manifest-level pruner reach any column with per-file statistics. let files_to_prune = pruning_predicate.prune(&PruneDataFiles::new( &schema, - &partition_schema, + &arrow_schema, &data_files, ))?; @@ -3003,6 +3012,117 @@ mod tests { .await; } + #[tokio::test] + pub async fn test_identity_self_named_partition_filter_prunes_files() { + // Regression for `partition_schema` being passed to `PruneDataFiles` at + // `table.rs:601`: identity-self-named partition columns (where the + // partition field's `name()` equals its `source_name()`) are dropped + // from `file_partition_fields` upstream, so `partition_schema` doesn't + // contain them. When a filter references such a column, the second- + // stage `PruneDataFiles` pruner fails its arrow-schema lookup in + // `min_values` / `max_values` and returns `None`, so no file gets + // pruned. The fix is to pass the full `arrow_schema` — which contains + // every column in the table — to `PruneDataFiles::new`. + // + // Reproducer: partition by `identity(kind)` on a string column named + // `kind`, insert rows for 3 distinct `kind` values (one parquet file + // per partition), then scan with a filter `kind = 'a'`. The resulting + // plan's file_groups should contain exactly ONE parquet file, not 3. + + use datafusion::physical_plan::displayable; + use datafusion::prelude::{col, lit}; + use datafusion::catalog::TableProvider; + + let object_store = ObjectStoreBuilder::memory(); + let catalog: Arc = Arc::new( + SqlCatalog::new("sqlite://", "identity_prune_probe", object_store) + .await + .unwrap(), + ); + + let schema = Schema::builder() + .with_struct_field(StructField { + id: 1, + name: "id".to_string(), + required: true, + field_type: Type::Primitive(PrimitiveType::Long), + doc: None, + }) + .with_struct_field(StructField { + id: 2, + name: "kind".to_string(), + required: true, + field_type: Type::Primitive(PrimitiveType::String), + doc: None, + }) + .build() + .unwrap(); + + // Identity-self-named: partition field "kind" on source column "kind". + let partition_spec = PartitionSpec::builder() + .with_partition_field(PartitionField::new(2, 1000, "kind", Transform::Identity)) + .build() + .expect("Failed to build partition spec"); + + let table = Table::builder() + .with_name("identity_prune_probe") + .with_location("/test/identity_prune_probe") + .with_schema(schema) + .with_partition_spec(partition_spec) + .build(&["test".to_owned()], catalog) + .await + .expect("Failed to create partitioned table"); + + let table = Arc::new(DataFusionTable::from(table)); + + let ctx = SessionContext::new(); + ctx.register_table("identity_prune_probe", table.clone()) + .unwrap(); + + // Three rows with three distinct kind values → three partition files. + ctx.sql( + "INSERT INTO identity_prune_probe (id, kind) VALUES + (1, 'a'), + (2, 'b'), + (3, 'c');", + ) + .await + .expect("Failed to create query plan for insert") + .collect() + .await + .expect("Failed to insert values into partitioned table"); + + // Sanity: three partition files exist unfiltered. + let state = ctx.state(); + let unfiltered_plan = table + .scan(&state, None, &[], None) + .await + .expect("unfiltered scan should succeed"); + let unfiltered_display = displayable(unfiltered_plan.as_ref()) + .indent(false) + .to_string(); + let unfiltered_parquet_count = unfiltered_display.matches(".parquet").count(); + assert_eq!( + unfiltered_parquet_count, 3, + "precondition: unfiltered scan should list all 3 partition files, got {unfiltered_parquet_count}:\n{unfiltered_display}" + ); + + // Now scan with a filter that matches exactly one partition. + let filter = col("kind").eq(lit("a")); + let filtered_plan = table + .scan(&state, None, &[filter], None) + .await + .expect("filtered scan should succeed"); + let filtered_display = displayable(filtered_plan.as_ref()) + .indent(false) + .to_string(); + let filtered_parquet_count = filtered_display.matches(".parquet").count(); + assert_eq!( + filtered_parquet_count, 1, + "expected pruning filter `kind = 'a'` to reduce scan to exactly 1 parquet file, got {filtered_parquet_count}:\n{filtered_display}" + ); + } + #[test] fn test_fake_object_store_url() { assert_eq!( From e7b9d0661f17a3795ad490b7ee618d1eb4115f84 Mon Sep 17 00:00:00 2001 From: Sergei Turukin Date: Wed, 15 Apr 2026 14:47:31 -0500 Subject: [PATCH 6/6] fix(manifest): resolve nested column ids in DataFile statistics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DataFile statistics maps (`lower_bounds`, `upper_bounds`, `column_sizes`, `value_counts`, `null_value_counts`, `nan_value_counts`) are keyed by global column id and Iceberg assigns those ids from the same pool at every depth — a struct field nested inside a list or inside a context/unstruct top-level column is just as valid a key as a top-level column. `AvroMap::into_value_map` was looking keys up via `StructType::get`, which only consults the *top-level* `lookup` table, so any nested id (e.g. Snowplow's `contexts_com_snowplowanalytics_*.*` fields reaching into the 400-700 range) silently failed with `ColumnNotInSchema`. That error was then `.unwrap()`'d inside `from_existing_with_filter`'s per-entry closure, which panicked the tokio worker and aborted the whole Lambda (`signal: aborted`) on any MERGE that touched a real Snowplow events table. Three fixes, smallest-to-largest: 1. `StructType::field_by_id(id)` — new recursive id lookup that walks nested `Struct`, `List`, and `Map` types. Independent from the existing top-level-only `get` so current callers of `get` are unaffected. 2. `AvroMap::into_value_map` now resolves ids via `field_by_id`. Unknown ids — entries pointing at fields that have been removed from the schema since the manifest was written — are now skipped rather than raised as `ColumnNotInSchema`. This matches Iceberg's schema-evolution semantics (old stats on removed fields are tolerated on read). 3. `iceberg-rust/src/table/manifest.rs::from_existing_with_filter`'s main rewrite loop is switched from `filter_map(...).unwrap()` to an explicit `for` loop that propagates per-entry errors via `?`. Any future deserialization edge case surfaces as a clean `Error` instead of a SIGABRT inside a tokio worker. Two new regression tests: - `types::tests::field_by_id_finds_nested_fields` — covers top-level, struct-of-struct, list, map, and unknown ids. - `manifest::tests::into_value_map_accepts_nested_field_ids` — builds an `AvroMap` with a nested-field key (479 inside a list), a top-level key, and an unknown key, and asserts all three paths (decode nested, decode top-level, silently skip unknown). Reproduced end-to-end: pre-fix, `MERGE INTO demo.atomic.events_hooli` aborts the Lambda after ~21s with `panicked at iceberg-rust/src/table/manifest.rs:549:18: ... Column 479 not in schema`. Co-Authored-By: Claude Opus 4.6 (1M context) --- iceberg-rust-spec/src/spec/manifest.rs | 119 +++++++++++++++---- iceberg-rust-spec/src/spec/types.rs | 151 +++++++++++++++++++++++++ iceberg-rust/src/table/manifest.rs | 23 ++-- 3 files changed, 260 insertions(+), 33 deletions(-) diff --git a/iceberg-rust-spec/src/spec/manifest.rs b/iceberg-rust-spec/src/spec/manifest.rs index 300f3212..3dfa4896 100644 --- a/iceberg-rust-spec/src/spec/manifest.rs +++ b/iceberg-rust-spec/src/spec/manifest.rs @@ -476,34 +476,30 @@ impl<'de, T: Serialize + DeserializeOwned + Clone> Deserialize<'de> for AvroMap< } impl AvroMap { - /// Converts a map of byte buffers into a map of typed Iceberg values using the provided schema. + /// Converts a map of byte buffers into a map of typed Iceberg values using + /// the provided schema. /// - /// # Arguments - /// * `schema` - The struct type schema used to determine the correct type for each value + /// DataFile statistics maps (`lower_bounds`, `upper_bounds`, ...) are keyed + /// by column id at any depth in the schema — including nested fields + /// inside struct/list/map types whose ids can be well above the top-level + /// count (common in Snowplow-style tables with context/unstruct arrays of + /// struct). Field lookup therefore walks the whole schema via + /// [`StructType::field_by_id`] rather than inspecting only top-level ids. /// - /// # Returns - /// * `Result, Error>` - A map of field IDs to their typed values, or an error if conversion fails + /// Ids that cannot be resolved anywhere in the schema are skipped rather + /// than treated as an error: the Iceberg spec permits stats for fields + /// that have since been removed from the schema, and a stale entry in an + /// old manifest must not break readers on any MERGE that rewrites it. fn into_value_map(self, schema: &StructType) -> Result, Error> { - Ok(HashMap::from_iter( - self.0 - .into_iter() - .map(|(k, v)| { - Ok(( - k, - Value::try_from_bytes( - &v, - &schema - .get(k as usize) - .ok_or(Error::ColumnNotInSchema( - k.to_string(), - format!("{schema:?}"), - ))? - .field_type, - )?, - )) - }) - .collect::, Error>>()?, - )) + self.0 + .into_iter() + .filter_map(|(k, v)| match schema.field_by_id(k) { + Some(field) => { + Some(Value::try_from_bytes(&v, &field.field_type).map(|val| (k, val))) + } + None => None, + }) + .collect::, Error>>() } } @@ -1582,4 +1578,77 @@ mod tests { assert_eq!(partition_values, result); } } + + /// Regression for the panic at `iceberg-rust/src/table/manifest.rs:549` + /// when MERGEing into real Snowplow events_hooli on S3 Tables. DataFile + /// statistics (`lower_bounds`, `upper_bounds`, ...) are keyed by column + /// id at any depth — including nested fields inside context/unstruct + /// structs whose ids can reach into the hundreds. `into_value_map` used + /// to validate these keys against only the top-level `StructType` + /// lookup, so any nested id returned `ColumnNotInSchema` and the + /// deserializer aborted. It now walks the schema recursively via + /// `StructType::field_by_id`, and tolerates ids that aren't in the + /// schema at all (schema evolution — Iceberg allows stats for removed + /// fields to linger in old manifests). + #[test] + fn into_value_map_accepts_nested_field_ids() { + use crate::spec::types::{ListType, StructType}; + use serde_bytes::ByteBuf; + + // Schema: top-level `event_id` (id=1) + `ctx_product` list + // whose element has an id-479 field. Mirrors the shape of a + // Snowplow `contexts_*` array. + let element_struct = Type::Struct(StructType::new(vec![ + StructField { + id: 479, + name: "product_id".to_string(), + required: false, + field_type: Type::Primitive(PrimitiveType::Int), + doc: None, + }, + ])); + let schema = StructType::new(vec![ + StructField { + id: 1, + name: "event_id".to_string(), + required: true, + field_type: Type::Primitive(PrimitiveType::String), + doc: None, + }, + StructField { + id: 2, + name: "ctx_product".to_string(), + required: false, + field_type: Type::List(ListType { + element_id: 478, + element_required: false, + element: Box::new(element_struct), + }), + doc: None, + }, + ]); + + // Build an AvroMap with a stat for the nested field id. + let mut raw: HashMap = HashMap::new(); + raw.insert(479, Value::Int(42).into()); + // Plus a known-unknown id to cover the schema-evolution case. + raw.insert(9999, Value::Int(7).into()); + // And a top-level id, to make sure we didn't regress that path. + raw.insert(1, Value::String("abc".to_string()).into()); + + let avro_map = AvroMap(raw); + let value_map = avro_map + .into_value_map(&schema) + .expect("nested field ids must deserialize cleanly"); + + // Nested lookup succeeded and decoded the int. + assert_eq!(value_map.get(&479), Some(&Value::Int(42))); + // Top-level lookup still works. + assert_eq!( + value_map.get(&1), + Some(&Value::String("abc".to_string())) + ); + // Unknown id is silently dropped, not an error. + assert!(value_map.get(&9999).is_none()); + } } diff --git a/iceberg-rust-spec/src/spec/types.rs b/iceberg-rust-spec/src/spec/types.rs index 0ae821f5..334c737c 100644 --- a/iceberg-rust-spec/src/spec/types.rs +++ b/iceberg-rust-spec/src/spec/types.rs @@ -300,6 +300,40 @@ impl StructType { .map(|idx| &self.fields[*idx]) } + /// Recursively looks up a `StructField` by its global field id. + /// + /// Iceberg assigns every struct field a globally unique id drawn from the + /// same pool regardless of depth, and DataFile statistics maps + /// (`lower_bounds`, `upper_bounds`, `column_sizes`, ...) are keyed by those + /// ids at any depth. This method walks the schema through nested `Struct`, + /// `List`, and `Map` types until it finds the field with the given id, or + /// returns `None` if no such field exists anywhere in the subtree. + pub fn field_by_id(&self, id: i32) -> Option<&StructField> { + if let Some(field) = self.get(id as usize) { + return Some(field); + } + for field in &self.fields { + if let Some(found) = Self::find_in_type(&field.field_type, id) { + return Some(found); + } + } + None + } + + fn find_in_type(ty: &Type, id: i32) -> Option<&StructField> { + match ty { + Type::Primitive(_) => None, + Type::Struct(s) => s.field_by_id(id), + Type::List(list) => Self::find_in_type(&list.element, id), + Type::Map(map) => { + if let Some(found) = Self::find_in_type(&map.key, id) { + return Some(found); + } + Self::find_in_type(&map.value, id) + } + } + } + /// Gets a reference to the StructField with the given name /// /// # Arguments @@ -633,4 +667,121 @@ mod tests { assert_eq!(Type::Primitive(PrimitiveType::String), *result.key); assert_eq!(Type::Primitive(PrimitiveType::Double), *result.value); } + + #[test] + fn field_by_id_finds_nested_fields() { + // Iceberg's DataFile statistics maps (lower_bounds, upper_bounds, + // column_sizes, ...) are keyed by column id, and those column ids + // can belong to fields at any depth — including fields nested + // under struct-of-struct, list, map<*,struct>. A caller + // that only inspects top-level ids will miss them and, in the + // manifest deserialization path, panic at `manifest.rs:549` on + // a `ColumnNotInSchema` error. Regression for that path. + let inner_struct = Type::Struct(StructType::new(vec![ + StructField { + id: 655, + name: "percent_progress".to_string(), + required: false, + field_type: Type::Primitive(PrimitiveType::Int), + doc: None, + }, + ])); + + let element_struct = Type::Struct(StructType::new(vec![ + StructField { + id: 479, + name: "product_id".to_string(), + required: false, + field_type: Type::Primitive(PrimitiveType::String), + doc: None, + }, + StructField { + id: 480, + name: "category".to_string(), + required: false, + field_type: Type::Primitive(PrimitiveType::String), + doc: None, + }, + ])); + + let map_value_struct = Type::Struct(StructType::new(vec![StructField { + id: 777, + name: "count".to_string(), + required: false, + field_type: Type::Primitive(PrimitiveType::Long), + doc: None, + }])); + + let outer = StructType::new(vec![ + StructField { + id: 1, + name: "event_id".to_string(), + required: true, + field_type: Type::Primitive(PrimitiveType::String), + doc: None, + }, + StructField { + id: 2, + name: "unstruct_event_media_ad_click".to_string(), + required: false, + field_type: inner_struct, + doc: None, + }, + StructField { + id: 3, + name: "contexts_ecommerce_product".to_string(), + required: false, + field_type: Type::List(ListType { + element_id: 478, + element_required: false, + element: Box::new(element_struct), + }), + doc: None, + }, + StructField { + id: 4, + name: "refr_params".to_string(), + required: false, + field_type: Type::Map(MapType { + key_id: 900, + key: Box::new(Type::Primitive(PrimitiveType::String)), + value_id: 901, + value_required: false, + value: Box::new(map_value_struct), + }), + doc: None, + }, + ]); + + // Top-level still works. + assert_eq!( + outer.field_by_id(1).map(|f| f.name.as_str()), + Some("event_id") + ); + + // Nested inside struct-of-struct. + assert_eq!( + outer.field_by_id(655).map(|f| f.name.as_str()), + Some("percent_progress") + ); + + // Nested inside list. + assert_eq!( + outer.field_by_id(479).map(|f| f.name.as_str()), + Some("product_id") + ); + assert_eq!( + outer.field_by_id(480).map(|f| f.name.as_str()), + Some("category") + ); + + // Nested inside map. + assert_eq!( + outer.field_by_id(777).map(|f| f.name.as_str()), + Some("count") + ); + + // Unknown id returns None (not panic, not a match). + assert!(outer.field_by_id(9999).is_none()); + } } diff --git a/iceberg-rust/src/table/manifest.rs b/iceberg-rust/src/table/manifest.rs index dafa5868..b2b46ea3 100644 --- a/iceberg-rust/src/table/manifest.rs +++ b/iceberg-rust/src/table/manifest.rs @@ -543,13 +543,20 @@ impl<'schema, 'metadata> ManifestWriter<'schema, 'metadata> { }, )?; - writer.extend(manifest_reader.filter_map(|entry| { - let mut entry = entry - .map_err(|err| apache_avro::Error::DeserializeValue(err.to_string())) - .unwrap(); + // Walk existing entries and either rewrite them as-is or mark them + // deleted if their file path is in the filter set. We collect the + // surviving values into a Vec because `writer.extend` cannot + // short-circuit on a per-element error the way a `for` loop can — a + // deserialization failure here used to be `.unwrap()`'d and took out + // the whole tokio worker (Lambda `signal: aborted`). + let mut rewritten = Vec::new(); + for entry in manifest_reader { + let mut entry = entry.map_err(|err| { + Error::from(apache_avro::Error::DeserializeValue(err.to_string())) + })?; if *entry.status() == Status::Deleted { - return None; + continue; } if entry.sequence_number().is_none() { @@ -567,12 +574,12 @@ impl<'schema, 'metadata> ManifestWriter<'schema, 'metadata> { filtered_stats.removed_data_files += 1; *entry.status_mut() = Status::Deleted; filtered_stats.filtered_entries.push(entry); - None } else { *entry.status_mut() = Status::Existing; - Some(to_value(entry).unwrap()) + rewritten.push(to_value(entry)?); } - }))?; + } + writer.extend(rewritten.into_iter())?; manifest.sequence_number = table_metadata.last_sequence_number + 1;