From 0dd80fb31cb66ca276475019dcc60227bc55955c Mon Sep 17 00:00:00 2001 From: JeonDaehong Date: Wed, 14 Oct 2026 21:22:36 +0900 Subject: [PATCH 1/2] fix(arrow): load equality deletes on fields nested in structs Equality ids may name a field nested in a struct, which the spec allows and EqualityDeleteWriter writes. Loading such a delete file failed because schema evolution only maps top-level field ids, so the delete file is now read by the top-level fields that hold the equality ids. The delete predicate also used the leaf name (`f`), which binds to a top-level column of the same name or fails; it now uses the full name (`s.f`). --- .../src/arrow/caching_delete_file_loader.rs | 104 ++++++++++++++++-- .../iceberg/src/arrow/delete_file_loader.rs | 28 ++++- 2 files changed, 123 insertions(+), 9 deletions(-) diff --git a/crates/iceberg/src/arrow/caching_delete_file_loader.rs b/crates/iceberg/src/arrow/caching_delete_file_loader.rs index 5582e4397f..529d90fed7 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, equality_delete_projection}; use crate::arrow::scan_metrics::ScanMetrics; use crate::arrow::{arrow_primitive_to_literal, arrow_schema_to_schema}; use crate::delete_vector::DeleteVector; @@ -302,8 +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 projection = equality_delete_projection(&schema, &equality_ids_vec); let evolved_stream = BasicDeleteFileLoader::evolve_schema( basic_delete_file_loader .parquet_to_batch_stream( @@ -313,7 +314,7 @@ impl CachingDeleteFileLoader { ) .await?, schema, - &equality_ids_vec, + &projection, ) .await?; @@ -604,7 +605,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 +662,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 +726,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 +899,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 +1802,83 @@ mod tests { assert_eq!(err.kind(), ErrorKind::DataInvalid); assert!(err.message().contains("expected 2 from record_count")); } + + /// An equality delete on a field nested in a struct loads, and its predicate uses the full + /// field name. The table also has a top-level column with the same leaf name, which the + /// predicate must not refer to. + #[tokio::test] + async fn test_nested_field_equality_delete_predicate() { + use crate::arrow::delete_file_loader::equality_delete_projection; + use crate::expr::Bind; + use crate::spec::{Datum, StructType}; + + let tmp_dir = TempDir::new().unwrap(); + let table_location = tmp_dir.path().to_str().unwrap(); + + let table_schema = 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(), + ])), + ) + .into(), + ]) + .build() + .unwrap(), + ); + + // Delete file with `s.f = 7`. + 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!("{table_location}/nested-eq-delete.parquet"); + let mut writer = + ArrowWriter::try_new(File::create(&path).unwrap(), delete_schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + // Only the top-level struct holding field 3 is read. + assert_eq!(equality_delete_projection(&table_schema, &[3]), vec![2]); + + let loader = BasicDeleteFileLoader::new(FileIO::new_with_fs(), ScanMetrics::new()); + let stream = loader + .parquet_to_batch_stream(&path, std::fs::metadata(&path).unwrap().len(), None) + .await + .unwrap(); + let evolved = BasicDeleteFileLoader::evolve_schema(stream, table_schema.clone(), &[2]) + .await + .unwrap(); + let predicate = CachingDeleteFileLoader::parse_equality_deletes_record_batch_stream( + evolved, + HashSet::from([3]), + ) + .await + .unwrap(); + + let s_f = Reference::new("s.f"); + assert_eq!( + predicate, + s_f.clone().is_null().or(s_f.not_equal_to(Datum::long(7))) + ); + // It binds to the nested field, not to the top-level `f`. + let bound = predicate.bind(table_schema, true).unwrap(); + assert!(format!("{bound:?}").contains("id: 3")); + assert!(!format!("{bound:?}").contains("id: 1,")); + } } diff --git a/crates/iceberg/src/arrow/delete_file_loader.rs b/crates/iceberg/src/arrow/delete_file_loader.rs index ec61710a42..8d628887f7 100644 --- a/crates/iceberg/src/arrow/delete_file_loader.rs +++ b/crates/iceberg/src/arrow/delete_file_loader.rs @@ -115,6 +115,30 @@ impl BasicDeleteFileLoader { } } +/// Returns the top-level field ids to read for an equality delete file. +/// +/// Equality ids may name fields nested in structs, but schema evolution of a delete file works on +/// top-level fields, so each id is replaced by the top-level field that contains it. +pub(crate) fn equality_delete_projection(schema: &Schema, equality_ids: &[i32]) -> Vec { + use crate::spec::{NestedField, Type}; + + fn contains(field: &NestedField, id: i32) -> bool { + field.id == id + || match field.field_type.as_ref() { + Type::Struct(s) => s.fields().iter().any(|child| contains(child, id)), + _ => false, + } + } + + schema + .as_struct() + .fields() + .iter() + .filter(|field| equality_ids.iter().any(|&id| contains(field, id))) + .map(|field| field.id) + .collect() +} + #[async_trait::async_trait] impl DeleteFileLoader for BasicDeleteFileLoader { async fn read_delete_file( @@ -130,10 +154,10 @@ 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(), + Some(ids) => equality_delete_projection(&schema, ids), None => schema.field_id_to_name_map().keys().cloned().collect(), }; From 43d31692ce07a5064de999fbe525e0969247c750 Mon Sep 17 00:00:00 2001 From: JeonDaehong Date: Wed, 28 Oct 2026 20:14:45 +0900 Subject: [PATCH 2/2] fix(arrow): read equality deletes through the schema pruned to their fields Equality ids missing from the scan schema now fail instead of being dropped. The delete file is read as the scan schema pruned to the equality fields, like Java's DeleteFilter, and a nested field that EqualityDeleteWriter wrote as a top-level column is put back under its parent struct first, so both layouts load. --- .../src/arrow/caching_delete_file_loader.rs | 217 ++++++++++++++---- .../iceberg/src/arrow/delete_file_loader.rs | 157 +++++++++++-- 2 files changed, 307 insertions(+), 67 deletions(-) diff --git a/crates/iceberg/src/arrow/caching_delete_file_loader.rs b/crates/iceberg/src/arrow/caching_delete_file_loader.rs index 529d90fed7..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, equality_delete_projection}; +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; @@ -304,8 +304,7 @@ impl CachingDeleteFileLoader { // Per the Iceberg spec, evolve schema for equality deletes but only for the // columns holding the equality_ids, not all table columns. let equality_ids_vec = task.equality_ids().unwrap().to_vec(); - let projection = equality_delete_projection(&schema, &equality_ids_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(), @@ -313,10 +312,9 @@ impl CachingDeleteFileLoader { task.key_metadata(), ) .await?, - schema, - &projection, - ) - .await?; + &schema, + &equality_ids_vec, + )?; Ok(DeleteFileContext::FreshEqDel { batch_stream: evolved_stream, @@ -1803,19 +1801,10 @@ mod tests { assert!(err.message().contains("expected 2 from record_count")); } - /// An equality delete on a field nested in a struct loads, and its predicate uses the full - /// field name. The table also has a top-level column with the same leaf name, which the - /// predicate must not refer to. - #[tokio::test] - async fn test_nested_field_equality_delete_predicate() { - use crate::arrow::delete_file_loader::equality_delete_projection; - use crate::expr::Bind; - use crate::spec::{Datum, StructType}; - - let tmp_dir = TempDir::new().unwrap(); - let table_location = tmp_dir.path().to_str().unwrap(); - - let table_schema = Arc::new( + /// 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(), @@ -1825,15 +1814,99 @@ mod tests { 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(); - // Delete file with `s.f = 7`. let f_field = Arc::new(simple_field("f", DataType::Int64, true, "3")); let s_field = simple_field( "s", @@ -1847,38 +1920,94 @@ mod tests { Arc::new(Int64Array::from(vec![7])) as ArrayRef, )]); let batch = RecordBatch::try_new(delete_schema.clone(), vec![Arc::new(s)]).unwrap(); - let path = format!("{table_location}/nested-eq-delete.parquet"); + 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(); - // Only the top-level struct holding field 3 is read. - assert_eq!(equality_delete_projection(&table_schema, &[3]), vec![2]); - - let loader = BasicDeleteFileLoader::new(FileIO::new_with_fs(), ScanMetrics::new()); - let stream = loader - .parquet_to_batch_stream(&path, std::fs::metadata(&path).unwrap().len(), None) + let predicate = load_eq_delete_predicate(&path, vec![3], schema.clone()) .await .unwrap(); - let evolved = BasicDeleteFileLoader::evolve_schema(stream, table_schema.clone(), &[2]) + 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(); - let predicate = CachingDeleteFileLoader::parse_equality_deletes_record_batch_stream( - evolved, - HashSet::from([3]), - ) - .await - .unwrap(); - let s_f = Reference::new("s.f"); - assert_eq!( - predicate, - s_f.clone().is_null().or(s_f.not_equal_to(Datum::long(7))) + // 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, ); - // It binds to the nested field, not to the top-level `f`. - let bound = predicate.bind(table_schema, true).unwrap(); - assert!(format!("{bound:?}").contains("id: 3")); - assert!(!format!("{bound:?}").contains("id: 1,")); + 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 8d628887f7..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,28 +119,134 @@ impl BasicDeleteFileLoader { } } -/// Returns the top-level field ids to read for an equality delete file. +/// 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. /// -/// Equality ids may name fields nested in structs, but schema evolution of a delete file works on -/// top-level fields, so each id is replaced by the top-level field that contains it. -pub(crate) fn equality_delete_projection(schema: &Schema, equality_ids: &[i32]) -> Vec { - use crate::spec::{NestedField, Type}; - - fn contains(field: &NestedField, id: i32) -> bool { - field.id == id - || match field.field_type.as_ref() { - Type::Struct(s) => s.fields().iter().any(|child| contains(child, id)), - _ => false, - } +/// 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) +} - schema +/// 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() - .filter(|field| equality_ids.iter().any(|&id| contains(field, id))) .map(|field| field.id) - .collect() + .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] @@ -156,12 +266,13 @@ impl DeleteFileLoader for BasicDeleteFileLoader { // 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) => equality_delete_projection(&schema, ids), - 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 + } + } } }