mbutrovich commented on code in PR #3453:
URL: https://github.com/apache/iceberg-rust/pull/3453#discussion_r4401021036
##########
crates/iceberg/src/arrow/caching_delete_file_loader.rs:
##########
@@ -302,20 +302,19 @@ 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(
Review Comment:
Should we check the equality ids against `schema` before
`try_start_eq_del_load` claims the file? Right now the new `DataInvalid` error
returns after `insert_equality_delete` has registered the receiver, so the
sender is dropped and the spawned task panics at
[`eq_del.await.unwrap()`](https://github.com/apache/iceberg-rust/blob/af1717ae9287c256b987d0418251443052ead66b/crates/iceberg/src/arrow/delete_filter.rs#L269).
Running `test_equality_delete_with_id_missing_from_schema_fails` on this head
prints:
```
panicked at crates/iceberg/src/arrow/delete_filter.rs:269:39:
called `Result::unwrap()` on an `Err` value: RecvError(())
```
That is the path in #3274: the entry stays `Loading`, and a later read of
the same delete file through the same `DeleteFilter` waits forever. #3274
covers I/O and parse failures in general, but this check needs only the scan
schema, so it can run before any state changes. It would also run before
`parquet_to_batch_stream` opens the file, which happens first today because
it's an argument to `evolve_equality_delete_stream`. Could the test also check
that a second `load_deletes` call on the same loader returns the error instead
of hanging?
##########
crates/iceberg/src/arrow/caching_delete_file_loader.rs:
##########
@@ -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<i32>,
+ schema: SchemaRef,
+ ) -> Result<Option<crate::expr::BoundPredicate>> {
+ 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<i32>) {
+ 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<crate::expr::BoundPredicate>,
+ 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() {
Review Comment:
What do you think about adding the end-to-end scan test in this PR, marked
`#[ignore = "needs #2906"]`? Then #2906 only has to remove the attribute, and
the test gets reviewed next to the loader code it covers. The repo already uses
a reasoned `#[ignore]` in
[`avro/schema.rs`](https://github.com/apache/iceberg-rust/blob/af1717ae9287c256b987d0418251443052ead66b/crates/iceberg/src/avro/schema.rs#L1021).
I'd cover both delete file layouts, a data row with a null `s`, a top-level
`f` with the same value, and a null `s` in the delete file on the required
`s.g`. The last one also needs #3306, so its ignore reason can name both PRs.
##########
crates/iceberg/src/arrow/delete_file_loader.rs:
##########
@@ -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<ArrowRecordBatchStream> {
+ 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<i32> = 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<RecordBatch> {
+ let top_level_ids: HashSet<i32> = 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);
+ }
Review Comment:
Could we work out the nesting once per stream instead of once per batch?
`top_level_ids`, `flat`, the output `Schema`, and the `flat.values().any(...)`
scan per column depend only on the batch schema and `pruned`, and the batch
schema is the same for every batch of a file. Computing the column mapping and
the output schema from the first batch (or from the Parquet file's Arrow
schema), then rebuilding each batch from the column indices, would drop the
per-batch `HashSet`, `HashMap`, and `Schema` allocations.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
