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..4ecda629 100644 --- a/datafusion_iceberg/src/table.rs +++ b/datafusion_iceberg/src/table.rs @@ -432,15 +432,88 @@ 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()); - // 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() @@ -453,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(); @@ -472,11 +537,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 @@ -529,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, ))?; @@ -633,13 +710,14 @@ 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; 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(); @@ -664,7 +742,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()) @@ -689,6 +770,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 +795,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 +824,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 +927,7 @@ async fn table_scan( last_updated_ms, enable_data_file_path_column, manifest_path, + &drop_partition_indices, ) }) .collect::, _>>()?; @@ -892,6 +977,7 @@ async fn table_scan( last_updated_ms, enable_data_file_path_column, manifest_path, + &drop_partition_indices, ) .unwrap() }) @@ -904,7 +990,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(); @@ -930,7 +1016,7 @@ async fn table_scan( } fn datafusion_partition_columns( - partition_fields: &[BoundPartitionField<'_>], + partition_fields: &[&BoundPartitionField<'_>], ) -> Result, DataFusionError> { let table_partition_cols: Vec = partition_fields .iter() @@ -1038,13 +1124,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)) @@ -2555,6 +2649,480 @@ 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; + } + + #[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!( 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/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"); + } } 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/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); + } } 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;