feat: use visitor to allow struct fields to be used in eq deletes

sdd · sdd · commit 681ab7e70b6c · 2025-09-01T07:34:40.000+01:00
diff --git a/crates/iceberg/src/arrow/caching_delete_file_loader.rs b/crates/iceberg/src/arrow/caching_delete_file_loader.rs
@@ -17,10 +17,10 @@
 
 use std::collections::{HashMap, HashSet};
 use std::ops::Not;
+use std::sync::Arc;
 
-use arrow_array::{Array, Int64Array, StringArray};
+use arrow_array::{Array, ArrayRef, Int64Array, StringArray, StructArray};
 use futures::{StreamExt, TryStreamExt};
-use itertools::Itertools;
 use tokio::sync::oneshot::{Receiver, channel};
 
 use super::delete_filter::DeleteFilter;
@@ -31,7 +31,11 @@ use crate::expr::Predicate::AlwaysTrue;
 use crate::expr::{Predicate, Reference};
 use crate::io::FileIO;
 use crate::scan::{ArrowRecordBatchStream, FileScanTaskDeleteFile};
-use crate::spec::{DataContentType, Datum, SchemaRef};
+use crate::spec::{
+    DataContentType, Datum, ListType, MapType, NestedField, NestedFieldRef, PartnerAccessor,
+    PrimitiveType, Schema, SchemaRef, SchemaWithPartnerVisitor, StructType, Type,
+    visit_schema_with_partner,
+};
 use crate::{Error, ErrorKind, Result};
 
 #[derive(Clone, Debug)]
@@ -321,6 +325,8 @@ impl CachingDeleteFileLoader {
         equality_ids: HashSet<i32>,
     ) -> Result<Predicate> {
         let mut result_predicate = AlwaysTrue;
+        let mut batch_schema_iceberg: Option<Schema> = None;
+        let accessor = EqDelRecordBatchPartnerAccessor;
 
         while let Some(record_batch) = stream.next().await {
             let record_batch = record_batch?;
@@ -329,47 +335,26 @@ impl CachingDeleteFileLoader {
                 return Ok(AlwaysTrue);
             }
 
-            let batch_schema_arrow = record_batch.schema();
-            let batch_schema_iceberg = arrow_schema_to_schema(batch_schema_arrow.as_ref())?;
-
-            let mut datum_columns_with_names: Vec<_> = record_batch
-                .columns()
-                .iter()
-                .zip(batch_schema_iceberg.as_struct().fields())
-                // only use columns that are in the set of equality_ids for this delete file
-                .filter(|(field, value)| equality_ids.contains(&value.id))
-                .map(|(column, field)| {
-                    let lit_vec = arrow_primitive_to_literal(column, &field.field_type)?;
-
-                    let primitive_type = field.field_type.as_primitive_type().ok_or(Error::new(
-                        ErrorKind::Unexpected,
-                        "field is not a primitive type",
-                    ))?;
-
-                    let datum_iterator: Box<dyn ExactSizeIterator<Item = Result<Option<Datum>>>> =
-                        Box::new(lit_vec.into_iter().map(move |c| {
-                            c.map(|literal| {
-                                literal
-                                    .as_primitive_literal()
-                                    .map(|primitive_literal| {
-                                        Datum::new(primitive_type.clone(), primitive_literal)
-                                    })
-                                    .ok_or(Error::new(
-                                        ErrorKind::Unexpected,
-                                        "failed to convert to primitive literal",
-                                    ))
-                            })
-                            .transpose()
-                        }));
-
-                    Ok::<_, Error>((datum_iterator, field.name.to_string()))
-                })
-                .try_collect()?;
+            let schema = match &batch_schema_iceberg {
+                Some(schema) => schema,
+                None => {
+                    let schema = arrow_schema_to_schema(record_batch.schema().as_ref())?;
+                    batch_schema_iceberg = Some(schema);
+                    batch_schema_iceberg.as_ref().unwrap()
+                }
+            };
+
+            let root_array: ArrayRef = Arc::new(StructArray::from(record_batch));
 
-            // consume all the iterators in lockstep, creating per-row predicates that get combined
-            // into a single final predicate
+            let mut processor = EqDelColumnProcessor::new(&equality_ids);
+            visit_schema_with_partner(schema, &root_array, &mut processor, &accessor)?;
 
-            // (2025-06-12) can't use `is_empty` as it depends on unstable library feature `exact_size_is_empty`
+            let mut datum_columns_with_names = processor.finish()?;
+            if datum_columns_with_names.is_empty() {
+                continue;
+            }
+
+            // Process the collected columns in lockstep
             #[allow(clippy::len_zero)]
             while datum_columns_with_names[0].0.len() > 0 {
                 let mut row_predicate = AlwaysTrue;
@@ -390,13 +375,169 @@ impl CachingDeleteFileLoader {
     }
 }
 
+struct EqDelColumnProcessor<'a> {
+    equality_ids: &'a HashSet<i32>,
+    collected_columns: Vec<(ArrayRef, String, Type)>,
+}
+
+impl<'a> EqDelColumnProcessor<'a> {
+    fn new(equality_ids: &'a HashSet<i32>) -> Self {
+        Self {
+            equality_ids,
+            collected_columns: Vec::with_capacity(equality_ids.len()),
+        }
+    }
+
+    #[allow(clippy::type_complexity)]
+    fn finish(
+        self,
+    ) -> Result<
+        Vec<(
+            Box<dyn ExactSizeIterator<Item = Result<Option<Datum>>>>,
+            String,
+        )>,
+    > {
+        self.collected_columns
+            .into_iter()
+            .map(|(array, field_name, field_type)| {
+                let primitive_type = field_type
+                    .as_primitive_type()
+                    .ok_or_else(|| {
+                        Error::new(ErrorKind::Unexpected, "field is not a primitive type")
+                    })?
+                    .clone();
+
+                let lit_vec = arrow_primitive_to_literal(&array, &field_type)?;
+                let datum_iterator: Box<dyn ExactSizeIterator<Item = Result<Option<Datum>>>> =
+                    Box::new(lit_vec.into_iter().map(move |c| {
+                        c.map(|literal| {
+                            literal
+                                .as_primitive_literal()
+                                .map(|primitive_literal| {
+                                    Datum::new(primitive_type.clone(), primitive_literal)
+                                })
+                                .ok_or(Error::new(
+                                    ErrorKind::Unexpected,
+                                    "failed to convert to primitive literal",
+                                ))
+                        })
+                        .transpose()
+                    }));
+
+                Ok((datum_iterator, field_name))
+            })
+            .collect::<Result<Vec<_>>>()
+    }
+}
+
+impl SchemaWithPartnerVisitor<ArrayRef> for EqDelColumnProcessor<'_> {
+    type T = ();
+
+    fn schema(&mut self, _schema: &Schema, _partner: &ArrayRef, _value: ()) -> Result<()> {
+        Ok(())
+    }
+
+    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() {
+            self.collected_columns.push((
+                partner.clone(),
+                field.name.clone(),
+                field.field_type.as_ref().clone(),
+            ));
+        }
+        Ok(())
+    }
+
+    fn r#struct(
+        &mut self,
+        _struct: &StructType,
+        _partner: &ArrayRef,
+        _results: Vec<()>,
+    ) -> Result<()> {
+        Ok(())
+    }
+
+    fn list(&mut self, _list: &ListType, _partner: &ArrayRef, _value: ()) -> Result<()> {
+        Ok(())
+    }
+
+    fn map(
+        &mut self,
+        _map: &MapType,
+        _partner: &ArrayRef,
+        _key_value: (),
+        _value: (),
+    ) -> Result<()> {
+        Ok(())
+    }
+
+    fn primitive(&mut self, _primitive: &PrimitiveType, _partner: &ArrayRef) -> Result<()> {
+        Ok(())
+    }
+}
+
+struct EqDelRecordBatchPartnerAccessor;
+
+impl PartnerAccessor<ArrayRef> for EqDelRecordBatchPartnerAccessor {
+    fn struct_partner<'a>(&self, schema_partner: &'a ArrayRef) -> Result<&'a ArrayRef> {
+        Ok(schema_partner)
+    }
+
+    fn field_partner<'a>(
+        &self,
+        struct_partner: &'a ArrayRef,
+        field: &NestedField,
+    ) -> Result<&'a ArrayRef> {
+        let Some(struct_array) = struct_partner.as_any().downcast_ref::<StructArray>() else {
+            return Err(Error::new(
+                ErrorKind::Unexpected,
+                "Expected struct array for field extraction",
+            ));
+        };
+
+        // Find the field by name within the struct
+        for (i, field_def) in struct_array.fields().iter().enumerate() {
+            if field_def.name() == &field.name {
+                return Ok(struct_array.column(i));
+            }
+        }
+
+        Err(Error::new(
+            ErrorKind::Unexpected,
+            format!("Field {} not found in parent struct", field.name),
+        ))
+    }
+
+    fn list_element_partner<'a>(&self, _list_partner: &'a ArrayRef) -> Result<&'a ArrayRef> {
+        Err(Error::new(
+            ErrorKind::FeatureUnsupported,
+            "List columns are unsupported in equality deletes",
+        ))
+    }
+
+    fn map_key_partner<'a>(&self, _map_partner: &'a ArrayRef) -> Result<&'a ArrayRef> {
+        Err(Error::new(
+            ErrorKind::FeatureUnsupported,
+            "Map columns are unsupported in equality deletes",
+        ))
+    }
+
+    fn map_value_partner<'a>(&self, _map_partner: &'a ArrayRef) -> Result<&'a ArrayRef> {
+        Err(Error::new(
+            ErrorKind::FeatureUnsupported,
+            "Map columns are unsupported in equality deletes",
+        ))
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::collections::HashMap;
     use std::fs::File;
     use std::sync::Arc;
 
-    use arrow_array::{ArrayRef, Int64Array, RecordBatch, StringArray};
+    use arrow_array::{ArrayRef, Int32Array, Int64Array, RecordBatch, StringArray, StructArray};
+    use arrow_schema::{DataType, Field, Fields};
     use parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY};
     use parquet::basic::Compression;
     use parquet::file::properties::WriterProperties;
@@ -419,7 +560,7 @@ mod tests {
             .await
             .expect("could not get batch stream");
 
-        let eq_ids = HashSet::from_iter(vec![2, 3, 4]);
+        let eq_ids = HashSet::from_iter(vec![2, 3, 4, 6]);
 
         let parsed_eq_delete = CachingDeleteFileLoader::parse_equality_deletes_record_batch_stream(
             record_batch_stream,
@@ -429,11 +570,19 @@ mod tests {
         .expect("error parsing batch stream");
         println!("{}", parsed_eq_delete);
 
-        let expected = "(((y != 1) OR (z != 100)) OR (a != \"HELP\")) AND (((y != 2) OR (z IS NOT NULL)) OR (a IS NOT NULL))".to_string();
+        let expected = "((((y != 1) OR (z != 100)) OR (a != \"HELP\")) OR (sa != 4)) AND ((((y != 2) OR (z IS NOT NULL)) OR (a IS NOT NULL)) OR (sa != 5))".to_string();
 
         assert_eq!(parsed_eq_delete.to_string(), expected);
     }
 
+    /// Create a simple field with metadata.
+    fn simple_field(name: &str, ty: DataType, nullable: bool, value: &str) -> Field {
+        arrow_schema::Field::new(name, ty, nullable).with_metadata(HashMap::from([(
+            PARQUET_FIELD_ID_META_KEY.to_string(),
+            value.to_string(),
+        )]))
+    }
+
     fn setup_write_equality_delete_file_1(table_location: &str) -> String {
         let col_y_vals = vec![1, 2];
         let col_y = Arc::new(Int64Array::from(col_y_vals)) as ArrayRef;
@@ -444,24 +593,42 @@ mod tests {
         let col_a_vals = vec![Some("HELP"), None];
         let col_a = Arc::new(StringArray::from(col_a_vals)) as ArrayRef;
 
+        let col_s = Arc::new(StructArray::from(vec![
+            (
+                Arc::new(simple_field("sa", DataType::Int32, false, "6")),
+                Arc::new(Int32Array::from(vec![4, 5])) as ArrayRef,
+            ),
+            (
+                Arc::new(simple_field("sb", DataType::Utf8, true, "7")),
+                Arc::new(StringArray::from(vec![Some("x"), None])) as ArrayRef,
+            ),
+        ]));
+
         let equality_delete_schema = {
+            let struct_field = DataType::Struct(Fields::from(vec![
+                simple_field("sa", DataType::Int32, false, "6"),
+                simple_field("sb", DataType::Utf8, true, "7"),
+            ]));
+
             let fields = vec![
-                arrow_schema::Field::new("y", arrow_schema::DataType::Int64, true).with_metadata(
-                    HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "2".to_string())]),
-                ),
-                arrow_schema::Field::new("z", arrow_schema::DataType::Int64, true).with_metadata(
-                    HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "3".to_string())]),
-                ),
-                arrow_schema::Field::new("a", arrow_schema::DataType::Utf8, true).with_metadata(
-                    HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "4".to_string())]),
-                ),
+                Field::new("y", arrow_schema::DataType::Int64, true).with_metadata(HashMap::from(
+                    [(PARQUET_FIELD_ID_META_KEY.to_string(), "2".to_string())],
+                )),
+                Field::new("z", arrow_schema::DataType::Int64, true).with_metadata(HashMap::from(
+                    [(PARQUET_FIELD_ID_META_KEY.to_string(), "3".to_string())],
+                )),
+                Field::new("a", arrow_schema::DataType::Utf8, true).with_metadata(HashMap::from([
+                    (PARQUET_FIELD_ID_META_KEY.to_string(), "4".to_string()),
+                ])),
+                simple_field("s", struct_field, false, "5"),
             ];
             Arc::new(arrow_schema::Schema::new(fields))
         };
 
-        let equality_deletes_to_write =
-            RecordBatch::try_new(equality_delete_schema.clone(), vec![col_y, col_z, col_a])
-                .unwrap();
+        let equality_deletes_to_write = RecordBatch::try_new(equality_delete_schema.clone(), vec![
+            col_y, col_z, col_a, col_s,
+        ])
+        .unwrap();
 
         let path = format!("{}/equality-deletes-1.parquet", &table_location);
 
diff --git a/crates/iceberg/src/scan/context.rs b/crates/iceberg/src/scan/context.rs
@@ -191,7 +191,6 @@ impl PlanContext {
 
         // TODO: Ideally we could ditch this intermediate Vec as we return an iterator.
         let mut filtered_mfcs = vec![];
-
         for manifest_file in manifest_files {
             let tx = if manifest_file.content == ManifestContentType::Deletes {
                 delete_file_tx.clone()
