kosiew commented on code in PR #24394:
URL: https://github.com/apache/datafusion/pull/24394#discussion_r3819824350
##########
datafusion/common/src/nested_struct.rs:
##########
@@ -1703,3 +1704,217 @@ mod tests {
));
}
}
+
+/// Adapts a [`RecordBatch`] to a target [`SchemaRef`].
+///
+/// If `batch` already has the target schema, it is returned immediately.
+///
+/// If `batch` has columns whose data types differ from `target_schema` (e.g.
stricter
+/// nested struct or list nullabilities), this function verifies that each
target data
+/// type contains the incoming column data type (as verified by
[`arrow::datatypes::DataType::contains`])
+/// and transforms the metadata/types of differing columns to match
`target_schema`
+/// without copying primitive buffer data.
+///
+/// If `batch` has an incompatible column count or incompatible column data
types,
+/// an error is returned.
+pub fn adapt_batch_to_schema(
+ batch: RecordBatch,
+ target_schema: &SchemaRef,
+) -> Result<RecordBatch> {
+ if Arc::ptr_eq(batch.schema_ref(), target_schema)
+ || batch.schema().as_ref() == target_schema.as_ref()
+ {
+ return Ok(batch);
+ }
+
+ if batch.num_columns() != target_schema.fields().len() {
+ return _plan_err!(
+ "Batch schema does not conform to expected schema (column count
mismatch). Expected: {target_schema}, got: {}",
+ batch.schema()
+ );
+ }
+
+ let mut columns = Vec::with_capacity(batch.num_columns());
+ let mut needs_column_adaptation = false;
+ let cast_options = CastOptions::default();
+
+ for (target_field, col) in
target_schema.fields().iter().zip(batch.columns()) {
+ if target_field.data_type() != col.data_type() {
+ // If data types differ, verify that target_field's data type
contains
+ // the column's data type (e.g. stricter nested struct / list
field nullability).
+ if !target_field.data_type().contains(col.data_type()) {
+ return _plan_err!(
+ "Batch column '{}' with type {} cannot be adapted to
expected type {}",
+ target_field.name(),
+ col.data_type(),
+ target_field.data_type()
+ );
+ }
+ needs_column_adaptation = true;
+ let adapted_col = cast_column(col, target_field.data_type(),
&cast_options)?;
Review Comment:
I think there is still an issue here for Union schemas.
`DataType::contains` accepts a target `Union` when its child field is
nullable and the corresponding incoming Union child is non-nullable. That means
the incoming schema is considered a valid stricter version of the target schema.
We then reach `cast_column`, but Arrow does not support casting a Union to
another Union. Its Union cast support is for extracting a child into a
non-Union target. In practice, this means a compatible stricter Union batch now
fails with `cannot cast Union ... to Union ...` instead of being adapted to the
declared input schema.
Could we handle Union arrays explicitly here by recursively adapting each
child to the corresponding target field, then rebuilding the `UnionArray` with
the target `UnionFields` while preserving the type IDs, dense offsets, and mode?
Another option would be to reject Union shapes that the adapter cannot
actually construct, even if `DataType::contains` currently considers them
compatible.
It would be good to add a direct `adapt_batch_to_schema` regression for
stricter Union child nullability, plus an aggregate execution regression that
exercises the same case through `AggregateExec`.
--
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]
