diff --git a/crates/iceberg/src/arrow/caching_delete_file_loader.rs b/crates/iceberg/src/arrow/caching_delete_file_loader.rs index 5582e4397f..800e86c115 100644 --- a/crates/iceberg/src/arrow/caching_delete_file_loader.rs +++ b/crates/iceberg/src/arrow/caching_delete_file_loader.rs @@ -24,7 +24,7 @@ use futures::{StreamExt, TryStreamExt}; use tokio::sync::oneshot::{Receiver, channel}; use super::delete_filter::{DeleteFilter, PosDelLoadAction}; -use crate::arrow::delete_file_loader::BasicDeleteFileLoader; +use crate::arrow::delete_file_loader::{BasicDeleteFileLoader, evolve_equality_delete_stream}; use crate::arrow::scan_metrics::ScanMetrics; use crate::arrow::{arrow_primitive_to_literal, arrow_schema_to_schema}; use crate::delete_vector::DeleteVector; @@ -302,9 +302,9 @@ impl CachingDeleteFileLoader { del_filter.insert_equality_delete(task.file_path(), receiver); // Per the Iceberg spec, evolve schema for equality deletes but only for the - // equality_ids columns, not all table columns. + // columns holding the equality_ids, not all table columns. let equality_ids_vec = task.equality_ids().unwrap().to_vec(); - let evolved_stream = BasicDeleteFileLoader::evolve_schema( + let evolved_stream = evolve_equality_delete_stream( basic_delete_file_loader .parquet_to_batch_stream( task.file_path(), @@ -312,10 +312,9 @@ impl CachingDeleteFileLoader { task.key_metadata(), ) .await?, - schema, + &schema, &equality_ids_vec, - ) - .await?; + )?; Ok(DeleteFileContext::FreshEqDel { batch_stream: evolved_stream, @@ -604,7 +603,7 @@ impl CachingDeleteFileLoader { let root_array: ArrayRef = Arc::new(StructArray::from(record_batch)); - let mut processor = EqDelColumnProcessor::new(&equality_ids); + let mut processor = EqDelColumnProcessor::new(&equality_ids, schema); visit_schema_with_partner(schema, &root_array, &mut processor, &accessor)?; let mut datum_columns_with_names = processor.finish()?; @@ -661,13 +660,15 @@ impl CachingDeleteFileLoader { struct EqDelColumnProcessor<'a> { equality_ids: &'a HashSet, + schema: &'a Schema, collected_columns: Vec<(ArrayRef, String, Type)>, } impl<'a> EqDelColumnProcessor<'a> { - fn new(equality_ids: &'a HashSet) -> Self { + fn new(equality_ids: &'a HashSet, schema: &'a Schema) -> Self { Self { equality_ids, + schema, collected_columns: Vec::with_capacity(equality_ids.len()), } } @@ -723,9 +724,17 @@ impl SchemaWithPartnerVisitor for EqDelColumnProcessor<'_> { fn field(&mut self, field: &NestedFieldRef, partner: &ArrayRef, _value: ()) -> Result<()> { if self.equality_ids.contains(&field.id) && field.field_type.as_primitive_type().is_some() { + // Use the full name so a nested field (`s.f`) binds to itself, not to a top-level + // column with the same leaf name. + let name = self.schema.name_by_field_id(field.id).ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + format!("equality delete field {} not found in schema", field.id), + ) + })?; self.collected_columns.push(( partner.clone(), - field.name.clone(), + name.to_string(), field.field_type.as_ref().clone(), )); } @@ -888,7 +897,7 @@ mod tests { .await .expect("error parsing batch stream"); - let expected = "((((((y IS NULL) OR (y != 1)) OR ((z IS NULL) OR (z != 100))) OR ((a IS NULL) OR (a != \"HELP\"))) OR ((sa IS NULL) OR (sa != 4))) OR ((b IS NULL) OR (b != 62696E6172795F64617461))) AND ((((((y IS NULL) OR (y != 2)) OR (z IS NOT NULL)) OR (a IS NOT NULL)) OR ((sa IS NULL) OR (sa != 5))) OR (b IS NOT NULL))".to_string(); + let expected = "((((((y IS NULL) OR (y != 1)) OR ((z IS NULL) OR (z != 100))) OR ((a IS NULL) OR (a != \"HELP\"))) OR ((s.sa IS NULL) OR (s.sa != 4))) OR ((b IS NULL) OR (b != 62696E6172795F64617461))) AND ((((((y IS NULL) OR (y != 2)) OR (z IS NOT NULL)) OR (a IS NOT NULL)) OR ((s.sa IS NULL) OR (s.sa != 5))) OR (b IS NOT NULL))".to_string(); assert_eq!(parsed_eq_delete.to_string(), expected); } @@ -1791,4 +1800,214 @@ mod tests { assert_eq!(err.kind(), ErrorKind::DataInvalid); assert!(err.message().contains("expected 2 from record_count")); } + + /// Top-level `f` (id 1) and `s.f` (id 3, in struct `s` with id 2), so a predicate that + /// binds by leaf name would pick the wrong column. + fn nested_eq_delete_table_schema() -> SchemaRef { + Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::optional(1, "f", Type::Primitive(PrimitiveType::Long)).into(), + NestedField::optional( + 2, + "s", + Type::Struct(StructType::new(vec![ + NestedField::optional(3, "f", Type::Primitive(PrimitiveType::Long)) + .into(), + NestedField::required(4, "g", Type::Primitive(PrimitiveType::Long)) + .into(), + ])), + ) + .into(), + ]) + .build() + .unwrap(), + ) + } + + /// Loads one equality delete file with `equality_ids` through `load_deletes` and returns the + /// predicate the delete filter builds for a scan of `schema`. + async fn load_eq_delete_predicate( + path: &str, + equality_ids: Vec, + schema: SchemaRef, + ) -> Result> { + use crate::scan::FileScanTask; + + let eq_del = FileScanTaskDeleteFile::builder() + .with_file_path(path.to_string()) + .with_file_size_in_bytes(std::fs::metadata(path).unwrap().len()) + .with_file_type(DataContentType::EqualityDeletes) + .with_file_format(DataFileFormat::Parquet) + .with_partition_spec_id(0) + .with_equality_ids(Some(equality_ids)) + .build() + .unwrap(); + let task = FileScanTask::builder() + .with_file_size_in_bytes(0) + .with_start(0) + .with_length(0) + .with_data_file_path(format!("{path}.data")) + .with_data_file_format(DataFileFormat::Parquet) + .with_schema(schema.clone()) + .with_project_field_ids(vec![1, 2]) + .with_deletes(vec![eq_del]) + .with_case_sensitive(true) + .build() + .unwrap(); + let loader = CachingDeleteFileLoader::new(FileIO::new_with_fs(), 10, Runtime::current()); + let filter = loader + .load_deletes(task.deletes(), task.schema_ref()) + .await + .unwrap()?; + filter.build_equality_delete_predicate(&task).await + } + + /// Field ids of the references in a bound predicate. + fn bound_field_ids(predicate: &crate::expr::BoundPredicate, ids: &mut Vec) { + use crate::expr::BoundPredicate; + match predicate { + BoundPredicate::And(expr) | BoundPredicate::Or(expr) => { + for input in expr.inputs() { + bound_field_ids(input, ids); + } + } + BoundPredicate::Not(expr) => bound_field_ids(expr.inputs()[0], ids), + BoundPredicate::Unary(expr) => ids.push(expr.term().field().id), + BoundPredicate::Binary(expr) => ids.push(expr.term().field().id), + BoundPredicate::Set(expr) => ids.push(expr.term().field().id), + BoundPredicate::AlwaysTrue | BoundPredicate::AlwaysFalse => {} + } + } + + /// The predicate for deleting `s.f = 7`, bound to field 3 and not to the top-level `f`. + fn assert_keeps_rows_unless_nested_f_is_7( + predicate: Option, + schema: SchemaRef, + ) { + use crate::expr::Bind; + + let predicate = predicate.expect("an equality delete predicate"); + let s_f = Reference::new("s.f"); + let expected = s_f + .clone() + .is_null() + .or(s_f.not_equal_to(Datum::long(7))) + .bind(schema, true) + .unwrap(); + assert_eq!(predicate, expected); + let mut ids = Vec::new(); + bound_field_ids(&predicate, &mut ids); + assert_eq!(ids, vec![3, 3]); + } + + /// A delete file in Java's layout: `s.f` under its parent struct `s`. + #[tokio::test] + async fn test_nested_field_equality_delete_in_struct_layout() { + let tmp_dir = TempDir::new().unwrap(); + let schema = nested_eq_delete_table_schema(); + + let f_field = Arc::new(simple_field("f", DataType::Int64, true, "3")); + let s_field = simple_field( + "s", + DataType::Struct(Fields::from(vec![f_field.clone()])), + true, + "2", + ); + let delete_schema = Arc::new(arrow_schema::Schema::new(vec![s_field])); + let s = StructArray::from(vec![( + f_field, + Arc::new(Int64Array::from(vec![7])) as ArrayRef, + )]); + let batch = RecordBatch::try_new(delete_schema.clone(), vec![Arc::new(s)]).unwrap(); + let path = format!( + "{}/nested-eq-delete.parquet", + tmp_dir.path().to_str().unwrap() + ); + let mut writer = + ArrowWriter::try_new(File::create(&path).unwrap(), delete_schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let predicate = load_eq_delete_predicate(&path, vec![3], schema.clone()) + .await + .unwrap(); + assert_keeps_rows_unless_nested_f_is_7(predicate, schema); + } + + /// A delete file from `EqualityDeleteWriter`, which writes `s.f` as a top-level column. + #[tokio::test] + async fn test_nested_field_equality_delete_from_equality_delete_writer() { + use crate::arrow::schema_to_arrow_schema; + use crate::writer::base_writer::equality_delete_writer::{ + EqualityDeleteFileWriterBuilder, EqualityDeleteWriterConfig, + }; + use crate::writer::file_writer::ParquetWriterBuilder; + use crate::writer::file_writer::location_generator::{ + DefaultFileNameGenerator, DefaultLocationGenerator, + }; + use crate::writer::file_writer::rolling_writer::RollingFileWriterBuilder; + use crate::writer::{IcebergWriter, IcebergWriterBuilder}; + + let tmp_dir = TempDir::new().unwrap(); + let file_io = FileIO::new_with_fs(); + let schema = nested_eq_delete_table_schema(); + + let config = EqualityDeleteWriterConfig::new(vec![3], schema.clone()).unwrap(); + let delete_schema = arrow_schema_to_schema(config.projected_arrow_schema_ref()).unwrap(); + let rolling = RollingFileWriterBuilder::new_with_default_file_size( + ParquetWriterBuilder::new(WriterProperties::builder().build(), Arc::new(delete_schema)), + file_io.clone(), + DefaultLocationGenerator::with_data_location( + tmp_dir.path().to_str().unwrap().to_string(), + ), + DefaultFileNameGenerator::new("eq-del".to_string(), None, DataFileFormat::Parquet), + ); + let mut writer = EqualityDeleteFileWriterBuilder::new(rolling, config) + .build(None) + .await + .unwrap(); + + // Rows of the table: f = 1, s = {f: 7, g: 0}. + let arrow_schema = Arc::new(schema_to_arrow_schema(&schema).unwrap()); + let DataType::Struct(s_fields) = arrow_schema.field(1).data_type().clone() else { + unreachable!() + }; + let s = StructArray::new( + s_fields, + vec![ + Arc::new(Int64Array::from(vec![7])) as ArrayRef, + Arc::new(Int64Array::from(vec![0])) as ArrayRef, + ], + None, + ); + let rows = RecordBatch::try_new(arrow_schema, vec![ + Arc::new(Int64Array::from(vec![1])) as ArrayRef, + Arc::new(s) as ArrayRef, + ]) + .unwrap(); + writer.write(rows).await.unwrap(); + let data_file = writer.close().await.unwrap().pop().unwrap(); + + let predicate = load_eq_delete_predicate(data_file.file_path(), vec![3], schema.clone()) + .await + .unwrap(); + assert_keeps_rows_unless_nested_f_is_7(predicate, schema); + } + + /// An equality id the scan schema does not have is an error, not a column to skip. + #[tokio::test] + async fn test_equality_delete_with_id_missing_from_schema_fails() { + let tmp_dir = TempDir::new().unwrap(); + let table_location = tmp_dir.path().to_str().unwrap(); + // Columns y (id 2) and z (id 3). + let path = setup_write_equality_delete_file_1(table_location); + let schema = nested_eq_delete_table_schema(); + + let err = load_eq_delete_predicate(&path, vec![1, 9], schema) + .await + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + assert!(err.message().contains("field id 9"), "{err}"); + } } diff --git a/crates/iceberg/src/arrow/delete_file_loader.rs b/crates/iceberg/src/arrow/delete_file_loader.rs index ec61710a42..f0cadce4d2 100644 --- a/crates/iceberg/src/arrow/delete_file_loader.rs +++ b/crates/iceberg/src/arrow/delete_file_loader.rs @@ -15,18 +15,22 @@ // specific language governing permissions and limitations // under the License. +use std::collections::{HashMap, HashSet}; use std::sync::Arc; +use arrow_array::{ArrayRef, RecordBatch, RecordBatchOptions, StructArray}; +use arrow_schema::{FieldRef, Fields, Schema as ArrowSchema}; use futures::{StreamExt, TryStreamExt}; -use parquet::arrow::ParquetRecordBatchStreamBuilder; +use parquet::arrow::{PARQUET_FIELD_ID_META_KEY, ParquetRecordBatchStreamBuilder}; use crate::arrow::ArrowReader; use crate::arrow::reader::ParquetReadOptions; use crate::arrow::record_batch_transformer::RecordBatchTransformerBuilder; use crate::arrow::scan_metrics::ScanMetrics; +use crate::arrow::schema::try_get_field_id_from_metadata; use crate::io::FileIO; use crate::scan::{ArrowRecordBatchStream, FileScanTaskDeleteFile}; -use crate::spec::{Schema, SchemaRef}; +use crate::spec::{Schema, SchemaRef, prune_columns}; use crate::{Error, ErrorKind, Result}; /// Delete File Loader @@ -115,6 +119,136 @@ impl BasicDeleteFileLoader { } } +/// Reads an equality delete file as `schema` pruned to its equality fields, the way Java's +/// `DeleteFilter` selects them, so a field nested in a struct keeps its parent struct. +/// +/// Java writes a nested equality field under its parent structs, while `EqualityDeleteWriter` +/// writes it as a top-level column. Such a column is put back under its parents before schema +/// evolution, so both layouts read the same way. +pub(crate) fn evolve_equality_delete_stream( + record_batch_stream: ArrowRecordBatchStream, + schema: &Schema, + equality_ids: &[i32], +) -> Result { + if let Some(id) = equality_ids + .iter() + .find(|&&id| schema.field_by_id(id).is_none()) + { + return Err(Error::new( + ErrorKind::DataInvalid, + format!("Equality delete field id {id} is not in the scan schema"), + )); + } + let crate::spec::Type::Struct(pruned) = + prune_columns(schema, equality_ids.iter().copied(), false)? + else { + return Err(Error::new( + ErrorKind::Unexpected, + "pruning a schema did not return a struct", + )); + }; + let projection: Vec = pruned.fields().iter().map(|field| field.id).collect(); + let pruned = Arc::new( + Schema::builder() + .with_schema_id(schema.schema_id()) + .with_fields(pruned.fields().to_vec()) + .build()?, + ); + let mut transformer = RecordBatchTransformerBuilder::new(pruned.clone(), &projection).build(); + Ok(Box::pin(record_batch_stream.map(move |batch| { + batch + .and_then(|batch| nest_top_level_equality_columns(batch, &pruned)) + .and_then(|batch| transformer.process_record_batch(batch)) + })) as ArrowRecordBatchStream) +} + +/// Moves top-level columns whose field ids belong to fields nested in `schema` back under their +/// parent structs. A batch with no such column is returned as is. +fn nest_top_level_equality_columns(batch: RecordBatch, schema: &Schema) -> Result { + let top_level_ids: HashSet = schema + .as_struct() + .fields() + .iter() + .map(|field| field.id) + .collect(); + let mut flat = HashMap::new(); + for (index, field) in batch.schema().fields().iter().enumerate() { + if let Some(id) = try_get_field_id_from_metadata(field)? + && !top_level_ids.contains(&id) + && schema.field_by_id(id).is_some() + { + flat.insert(id, index); + } + } + if flat.is_empty() { + return Ok(batch); + } + + fn nest( + field: &crate::spec::NestedField, + batch: &RecordBatch, + flat: &HashMap, + ) -> Option<(FieldRef, ArrayRef)> { + if let Some(&index) = flat.get(&field.id) { + return Some(( + batch.schema().field(index).clone().into(), + batch.column(index).clone(), + )); + } + let crate::spec::Type::Struct(struct_type) = field.field_type.as_ref() else { + return None; + }; + let (fields, columns): (Vec, Vec) = struct_type + .fields() + .iter() + .filter_map(|child| nest(child, batch, flat)) + .unzip(); + if fields.is_empty() { + return None; + } + let fields = Fields::from(fields); + let array = StructArray::new(fields.clone(), columns, None); + let arrow_field = + arrow_schema::Field::new(&field.name, arrow_schema::DataType::Struct(fields), true) + .with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + field.id.to_string(), + )])); + Some((Arc::new(arrow_field), Arc::new(array))) + } + + let mut fields: Vec = Vec::new(); + let mut columns: Vec = Vec::new(); + for (index, field) in batch.schema().fields().iter().enumerate() { + if !flat.values().any(|&flat_index| flat_index == index) { + fields.push(field.clone()); + columns.push(batch.column(index).clone()); + } + } + let present: HashSet = fields + .iter() + .map(|field| try_get_field_id_from_metadata(field)) + .collect::>>()? + .into_iter() + .flatten() + .collect(); + for field in schema.as_struct().fields() { + if present.contains(&field.id) { + continue; + } + if let Some((arrow_field, column)) = nest(field, &batch, &flat) { + fields.push(arrow_field); + columns.push(column); + } + } + let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows())); + Ok(RecordBatch::try_new_with_options( + Arc::new(ArrowSchema::new(fields)), + columns, + &options, + )?) +} + #[async_trait::async_trait] impl DeleteFileLoader for BasicDeleteFileLoader { async fn read_delete_file( @@ -130,14 +264,15 @@ impl DeleteFileLoader for BasicDeleteFileLoader { ) .await?; - // For equality deletes, only evolve the equality_ids columns. + // For equality deletes, only evolve the columns holding the equality_ids. // For positional deletes (equality_ids is None), use all field IDs. - let field_ids = match task.equality_ids() { - Some(ids) => ids.to_vec(), - None => schema.field_id_to_name_map().keys().cloned().collect(), - }; - - Self::evolve_schema(raw_batch_stream, schema, &field_ids).await + match task.equality_ids() { + Some(ids) => evolve_equality_delete_stream(raw_batch_stream, &schema, ids), + None => { + let field_ids: Vec = schema.field_id_to_name_map().keys().cloned().collect(); + Self::evolve_schema(raw_batch_stream, schema, &field_ids).await + } + } } }