JeonDaehong commented on code in PR #3453:
URL: https://github.com/apache/iceberg-rust/pull/3453#discussion_r4537625719
##########
crates/iceberg/src/arrow/delete_file_loader.rs:
##########
@@ -115,6 +119,195 @@ impl BasicDeleteFileLoader {
}
}
+/// Checks that every equality id is a field of the scan schema. Java's
`DeleteFilter` fails the
+/// same way rather than matching rows on fewer columns.
+pub(crate) fn validate_equality_ids(schema: &Schema, equality_ids: &[i32]) ->
Result<()> {
+ match equality_ids
+ .iter()
+ .find(|&&id| schema.field_by_id(id).is_none())
+ {
+ Some(id) => Err(Error::new(
+ ErrorKind::DataInvalid,
+ format!("Equality delete field id {id} is not in the scan schema"),
+ )),
+ None => Ok(()),
+ }
+}
+
+/// 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> {
+ validate_equality_ids(schema, equality_ids)?;
+ 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();
+ // The batches of a file share its schema, so the nesting is worked out
once and only
+ // rebuilt if a batch comes with a different schema.
+ let mut nesting: Option<EqualityColumnNesting> = None;
+ Ok(Box::pin(record_batch_stream.map(move |batch| {
+ let batch = batch?;
+ let nesting = match &mut nesting {
+ Some(nesting) if nesting.input_schema == batch.schema() => nesting,
+ slot => slot.insert(EqualityColumnNesting::try_new(batch.schema(),
&pruned)?),
Review Comment:
Added `test_equality_delete_stream_renests_on_schema_change`: a flat, a
struct-layout, then a flat batch through one stream. It fails if the first
batch's nesting is reused, which `nests_every_batch` doesn't catch.
##########
crates/iceberg/src/arrow/delete_file_loader.rs:
##########
@@ -115,6 +119,195 @@ impl BasicDeleteFileLoader {
}
}
+/// Checks that every equality id is a field of the scan schema. Java's
`DeleteFilter` fails the
+/// same way rather than matching rows on fewer columns.
+pub(crate) fn validate_equality_ids(schema: &Schema, equality_ids: &[i32]) ->
Result<()> {
+ match equality_ids
+ .iter()
+ .find(|&&id| schema.field_by_id(id).is_none())
+ {
+ Some(id) => Err(Error::new(
+ ErrorKind::DataInvalid,
+ format!("Equality delete field id {id} is not in the scan schema"),
+ )),
+ None => Ok(()),
+ }
+}
+
+/// 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> {
+ validate_equality_ids(schema, equality_ids)?;
+ 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();
+ // The batches of a file share its schema, so the nesting is worked out
once and only
+ // rebuilt if a batch comes with a different schema.
+ let mut nesting: Option<EqualityColumnNesting> = None;
+ Ok(Box::pin(record_batch_stream.map(move |batch| {
+ let batch = batch?;
+ let nesting = match &mut nesting {
+ Some(nesting) if nesting.input_schema == batch.schema() => nesting,
+ slot => slot.insert(EqualityColumnNesting::try_new(batch.schema(),
&pruned)?),
+ };
+ transformer.process_record_batch(nesting.apply(batch)?)
+ })) as ArrowRecordBatchStream)
+}
+
+/// How to build the columns of an equality delete batch once the top-level
columns of fields
+/// nested in the pruned schema are moved back under their parent structs.
+struct EqualityColumnNesting {
+ input_schema: ArrowSchemaRef,
+ /// The output schema and columns, or `None` when no column has to move.
+ output: Option<(ArrowSchemaRef, Vec<NestedColumn>)>,
+}
+
+/// An output column: a column of the input batch, or a struct built from such
columns.
+enum NestedColumn {
+ Input(usize),
+ Struct(Fields, Vec<NestedColumn>),
+}
+
+impl NestedColumn {
+ fn build(&self, batch: &RecordBatch) -> ArrayRef {
+ match self {
+ NestedColumn::Input(index) => batch.column(*index).clone(),
+ NestedColumn::Struct(fields, children) =>
Arc::new(StructArray::new(
+ fields.clone(),
+ children.iter().map(|child| child.build(batch)).collect(),
+ None,
+ )),
+ }
+ }
+}
+
+impl EqualityColumnNesting {
+ fn try_new(input_schema: ArrowSchemaRef, schema: &Schema) -> Result<Self> {
+ 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 input_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(Self {
+ input_schema,
+ output: None,
+ });
+ }
+
+ fn nest(
Review Comment:
Added `[3,4]` from the writer and `a.b.c` in both layouts. In the
struct-layout one `c` is required, since an optional `c` already comes back
null from the Parquet reader; it fails if `push_struct_nulls_down` stops one
level down. The parent now takes its nullability from `field.required`. The
required-`s` case passes without that change too, because the transformer casts
the batch to the pruned schema, so it's a cleanup rather than a fix.
##########
crates/iceberg/src/arrow/caching_delete_file_loader.rs:
##########
@@ -619,9 +624,10 @@ impl CachingDeleteFileLoader {
}
};
- let root_array: ArrayRef =
Arc::new(StructArray::from(record_batch));
+ let root_array: ArrayRef =
+
push_struct_nulls_down(&(Arc::new(StructArray::from(record_batch)) as
ArrayRef))?;
Review Comment:
I'd keep the spec reading here. Telling a null `s` from `s = {f: null}`
needs a predicate on `s` itself rather than the leaf. Once #2906 makes this
reachable I can open an issue for the difference from Java.
--
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]
