patrickswedish commented on code in PR #24394:
URL: https://github.com/apache/datafusion/pull/24394#discussion_r3824166203
##########
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:
Hi @kosiew,
Thank you for the guidance! We have addressed this in the latest update:
1. **Narrowed Schema Conformance for Unions (cast_union_column &
alidate_union_schema_compatibility)**:
- Explicitly handles Union arrays (both Sparse and Dense modes) in
ested_struct::cast_column.
- Requires matching union modes, exact type ID set equality, and
recursive containment ( arget_child.contains(source_child)).
- Recursively adapts matching children using cast_column and reconstructs
the UnionArray with arget_fields, preserving ype_ids and dense
offsets buffers without copying buffer data.
- Rejects unsupported field set evolution (extra/missing type IDs).
- Removed Union from equires_nested_struct_cast so that generic SQL CAST
semantics across DataFusion are untouched.
2. **Unit and Integration Regressions**:
- Added unit tests verifying exact unpacked scalar values (10, "b", 30),
row-level ype_ids, offsets, and target schema containment for Dense and
Sparse unions.
- Added non-contiguous/reordered type-ID test ([(1, int), (3, str)] ->
[(3, str), (1, int)]).
- Added negative tests verifying rejection of field-set mismatches and
mode mismatches.
- Added direct AggregateExec input boundary integration tests in
ested_nullability.rs asserting unpacked rray_agg(b) values and output
schema nullability.
--
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]
