alamb commented on code in PR #24394:
URL: https://github.com/apache/datafusion/pull/24394#discussion_r3824066189
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -876,6 +877,45 @@ pub struct AggregateExec {
dynamic_filter: Option<Arc<AggrDynFilter>>,
}
+/// A stream wrapper that ensures every yielded batch matches the declared
input schema.
Review Comment:
I feel like this is a solution to a symptom (nullability mismatch) rather
than the underlying problem (an operator is not declaring its output schema
correctly and producing record batches with the wrong shape)
##########
datafusion/common/src/nested_struct.rs:
##########
@@ -1703,3 +1833,525 @@ 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)?;
+ columns.push(adapted_col);
+ } else {
+ columns.push(Arc::clone(col));
+ }
+ }
+
+ if needs_column_adaptation {
+ Ok(RecordBatch::try_new(Arc::clone(target_schema), columns)?)
+ } else {
+ // Schema differs only in top-level metadata or field nullability,
while
+ // column data types match exactly. Replace the schema on the batch.
+ Ok(RecordBatch::try_new(
+ Arc::clone(target_schema),
+ batch.columns().to_vec(),
+ )?)
+ }
+}
+
+#[cfg(test)]
+mod adapt_schema_tests {
Review Comment:
There seems to be a lot of repetition in these tests -- perhaps some of the
hints here could be applied:
https://datafusion.apache.org/contributor-guide/pr_review.html#review-the-test-coverage
--
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]
