diff --git a/crates/iceberg/src/arrow/record_batch_transformer.rs b/crates/iceberg/src/arrow/record_batch_transformer.rs index 76ea4fe28c..ec91eb7b3c 100644 --- a/crates/iceberg/src/arrow/record_batch_transformer.rs +++ b/crates/iceberg/src/arrow/record_batch_transformer.rs @@ -1484,6 +1484,83 @@ mod test { assert_eq!(get_string_value(result.column(2).as_ref(), 1), "Bob"); } + /// Test that a dropped identity partition source column is skipped rather than erroring. + /// + /// A historical spec may use `identity(col)` where `col` was later dropped from the + /// schema. `constants_map()` must skip that field instead of failing, so the file can + /// still be read (the dropped column is not projected). + #[test] + fn dropped_identity_partition_source_column_is_skipped() { + use crate::spec::{Struct, Transform}; + + // Original schema used to build the spec: has id, dept, name. + let original_schema = Arc::new( + Schema::builder() + .with_schema_id(0) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::required(2, "dept", Type::Primitive(PrimitiveType::String)).into(), + NestedField::optional(3, "name", Type::Primitive(PrimitiveType::String)).into(), + ]) + .build() + .unwrap(), + ); + + // Spec: identity(dept). Built against the original schema where dept still exists. + let partition_spec = Arc::new( + crate::spec::PartitionSpec::builder(original_schema.clone()) + .with_spec_id(0) + .add_partition_field("dept", "dept", Transform::Identity) + .unwrap() + .build() + .unwrap(), + ); + + // Evolved snapshot schema: dept (field id 2) has been dropped. + let evolved_schema = Arc::new( + Schema::builder() + .with_schema_id(1) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional(3, "name", Type::Primitive(PrimitiveType::String)).into(), + ]) + .build() + .unwrap(), + ); + + let partition_data = Struct::from_iter(vec![Some(Literal::string("engineering"))]); + + let parquet_schema = Arc::new(ArrowSchema::new(vec![ + simple_field("id", DataType::Int32, false, "1"), + simple_field("name", DataType::Utf8, true, "3"), + ])); + + let projected_field_ids = [1, 3]; // id, name -- dept is gone + + // The dropped source column must not cause an error here. + let mut transformer = + RecordBatchTransformerBuilder::new(evolved_schema, &projected_field_ids) + .with_partition(partition_spec, partition_data) + .expect("dropped partition source column should be skipped, not error") + .build(); + + let parquet_batch = RecordBatch::try_new(parquet_schema, vec![ + Arc::new(Int32Array::from(vec![100, 200])), + Arc::new(StringArray::from(vec!["Alice", "Bob"])), + ]) + .unwrap(); + + let result = transformer.process_record_batch(parquet_batch).unwrap(); + + // Only the live columns are projected; no constant is added for the dropped dept. + assert_eq!(result.num_columns(), 2); + assert_eq!(result.num_rows(), 2); + assert_eq!(get_int_value(result.column(0).as_ref(), 0), 100); + assert_eq!(get_int_value(result.column(0).as_ref(), 1), 200); + assert_eq!(get_string_value(result.column(1).as_ref(), 0), "Alice"); + assert_eq!(get_string_value(result.column(1).as_ref(), 1), "Bob"); + } + /// Test bucket partitioning with renamed source column. /// /// This verifies correct behavior for TestRuntimeFiltering.testRenamedSourceColumnTable() in Iceberg Java. diff --git a/crates/iceberg/src/scan/cache.rs b/crates/iceberg/src/scan/cache.rs index 66242e389b..b5eb3c4ce0 100644 --- a/crates/iceberg/src/scan/cache.rs +++ b/crates/iceberg/src/scan/cache.rs @@ -69,40 +69,65 @@ impl PartitionFilterCache { format!("Could not find partition spec for id {spec_id}"), ))?; - let partition_type = partition_spec.partition_type(schema)?; - let partition_fields = partition_type.fields().to_owned(); - let partition_schema = Arc::new( - Schema::builder() - .with_schema_id(partition_spec.spec_id()) - .with_fields(partition_fields) - .build()?, - ); + // A historical spec may reference a source column that was later dropped from the + // schema, which is a legitimate v2+ state. Such a spec cannot be resolved to a + // partition type, so it falls back to an always-true filter: files under the spec + // are not partition-pruned but still receive the row filter. Any other resolution + // failure is unexpected and propagates. The fallback is cached by spec id like any + // other filter; this is safe only because the cache lives per-scan in `PlanContext` + // with a fixed schema and predicate. Hoisting it to table or catalog scope would + // pin a spec to always-true even for a later scan whose schema could resolve it. + // TODO(https://github.com/apache/iceberg-rust/issues/2844): derive partition types from + // transforms where possible to restore pruning on a historical spec's still-live fields, + // which also makes the fallback per-term (dropped fields only) instead of per-spec. + let has_dropped_source_column = partition_spec + .fields() + .iter() + .any(|field| schema.field_by_id(field.source_id).is_none()); - let mut inclusive_projection = InclusiveProjection::new(partition_spec.clone()); + let partition_filter = if has_dropped_source_column { + BoundPredicate::AlwaysTrue + } else { + let partition_type = partition_spec.partition_type(schema)?; + let partition_fields = partition_type.fields().to_owned(); + let partition_schema = Arc::new( + Schema::builder() + .with_schema_id(partition_spec.spec_id()) + .with_fields(partition_fields) + .build()?, + ); - let partition_filter = inclusive_projection - .project(&filter)? - .rewrite_not() - .bind(partition_schema.clone(), case_sensitive)?; + let mut inclusive_projection = InclusiveProjection::new(partition_spec.clone()); - self.0 - .write() - .map_err(|_| { - Error::new( - ErrorKind::Unexpected, - "PartitionFilterCache RwLock was poisoned", - ) - })? - .insert(spec_id, Arc::new(partition_filter)); + inclusive_projection + .project(&filter)? + .rewrite_not() + .bind(partition_schema.clone(), case_sensitive)? + }; - let read = self.0.read().map_err(|_| { + let mut write = self.0.write().map_err(|_| { Error::new( ErrorKind::Unexpected, "PartitionFilterCache RwLock was poisoned", ) })?; - Ok(read.get(&spec_id).unwrap().clone()) + // Another thread may have populated the entry while the filter was computed without a + // lock held. Only warn and insert when this call is the one that populates it, so the + // always-true fallback is logged once per spec rather than once per racing thread. + let partition_filter = write.entry(spec_id).or_insert_with(|| { + if has_dropped_source_column { + tracing::warn!( + spec_id, + "Partition spec references a source column not in the scan schema; \ + skipping partition pruning for its manifests" + ); + } + + Arc::new(partition_filter) + }); + + Ok(partition_filter.clone()) } } @@ -235,3 +260,81 @@ impl ExpressionEvaluatorCache { Ok(read.get(&spec_id).unwrap().clone()) } } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use super::PartitionFilterCache; + use crate::expr::{Bind, BoundPredicate, Reference}; + use crate::spec::{ + Datum, FormatVersion, NestedField, PrimitiveType, Schema, SortOrder, TableMetadataBuilder, + Transform, Type, UnboundPartitionSpec, + }; + + /// A historical spec whose source column was dropped from the current schema resolves to an + /// always-true filter, and the fallback is cached under the spec id. + #[test] + fn dropped_partition_source_column_falls_back_to_always_true() { + let schema = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(), + NestedField::required(2, "part", Type::Primitive(PrimitiveType::Long)).into(), + ]) + .build() + .unwrap(); + + // Spec 0 partitions on `part` (id 2). + let spec = UnboundPartitionSpec::builder() + .with_spec_id(0) + .add_partition_field(2, "part", Transform::Identity) + .unwrap() + .build(); + + // Evolve the schema so that `part` is no longer present, leaving spec 0 unresolvable. + let evolved_schema = Schema::builder() + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(), + ]) + .build() + .unwrap(); + // Make an unpartitioned spec the default before dropping `part`, so spec 0 survives + // only as a historical spec that can no longer be resolved against the schema. + let table_metadata = Arc::new( + TableMetadataBuilder::new( + schema, + spec, + SortOrder::unsorted_order(), + "s3://bucket/test".to_string(), + FormatVersion::V2, + HashMap::new(), + ) + .unwrap() + .add_default_partition_spec(UnboundPartitionSpec::builder().build()) + .unwrap() + .add_current_schema(evolved_schema.clone()) + .unwrap() + .build() + .unwrap() + .metadata, + ); + + let filter = Reference::new("id") + .greater_than_or_equal_to(Datum::long(5)) + .bind(Arc::new(evolved_schema.clone()), true) + .unwrap(); + + let cache = PartitionFilterCache::new(); + let partition_filter = cache + .get(0, &table_metadata, &evolved_schema, true, filter.clone()) + .unwrap(); + assert!(matches!(*partition_filter, BoundPredicate::AlwaysTrue)); + + // The fallback is cached, so a second lookup returns the same instance. + let cached = cache + .get(0, &table_metadata, &evolved_schema, true, filter) + .unwrap(); + assert!(Arc::ptr_eq(&partition_filter, &cached)); + } +} diff --git a/crates/iceberg/src/scan/context.rs b/crates/iceberg/src/scan/context.rs index 4bd7e31de1..9181d02727 100644 --- a/crates/iceberg/src/scan/context.rs +++ b/crates/iceberg/src/scan/context.rs @@ -175,23 +175,11 @@ impl PlanContext { .await } - /// Returns the partition filter for a manifest, or an always-true filter when the - /// manifest's spec cannot be resolved against the snapshot schema. Files of such - /// manifests are not partition-pruned but still receive the row filter. + /// Returns the partition filter for a manifest. See [`PartitionFilterCache::get`] for the + /// always-true fallback when the manifest's spec cannot be resolved against the scan schema. fn get_partition_filter(&self, manifest_file: &ManifestFile) -> Result> { let partition_spec_id = manifest_file.partition_spec_id; - // Historical specs may reference source columns that were later dropped from the - // schema, in which case the partition type cannot be resolved. A missing spec is - // reported by the partition filter cache instead. - let resolvable = self - .table_metadata - .partition_spec_by_id(partition_spec_id) - .is_none_or(|spec| spec.partition_type(&self.snapshot_schema).is_ok()); - if !resolvable { - return Ok(Arc::new(BoundPredicate::AlwaysTrue)); - } - let partition_filter = self.partition_filter_cache.get( partition_spec_id, &self.table_metadata,