patrickswedish commented on code in PR #24394:
URL: https://github.com/apache/datafusion/pull/24394#discussion_r3824595366
##########
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:
Hi @alamb,
Thank you for this key architectural insight! You are completely right:
fixing this at the \AggregateExec\ boundary was treating a symptom rather than
addressing the producer contract.
### Root Cause Analysis
In DataFusion, in-memory table sources like \MemTable::try_new\ accept
batches whose schemas are stricter than the declared table schema using
\Schema::contains(&batches_schema)\ (e.g. nullable nested fields in the
declared table schema vs non-nullable nested fields in the input batches).
When \MemTable::scan\ creates \MemorySourceConfig\ / \MemoryExec\,
\MemoryStream\ was constructed with the declared schema, but its \poll_next\
emitted the underlying stricter \RecordBatch\s without adapting them.
Downstream operators (like \AggregateExec\) received batches that did not
conform to the stream's advertised output schema.
### Architectural Solution
1. **Reverted AggregateExec Changes**: Completely removed
\AdaptedInputRecordBatchStream\ and \AggregateExec::execute_input\, restoring
all aggregate physical plans to clean upstream state.
2. **Fixed Producer Invariant in \MemoryStream\
(\physical-plan/src/memory.rs\)**:
- \MemoryStream::poll_next\ now normalizes emitted batches using
\dapt_batch_to_schema(batch, &self.schema)\ whenever a batch differs from
\self.schema\ and \self.schema.contains(batch.schema())\.
- Every \RecordBatch\ emitted by \MemoryStream\ is guaranteed to conform
to \stream.schema()\.
3. **Retained Narrow Schema Conformance in \
ested_struct\ (\common/src/nested_struct.rs\)**:
- \dapt_batch_to_schema\ supports Structs, Lists, and Unions
(Dense/Sparse) without changing general DataFusion CAST behavior (\
equires_nested_struct_cast\ remains unchanged).
4. **Unit and Integration Regressions**:
- Added unit tests in \physical-plan/src/memory.rs\ directly verifying
that \MemoryStream\ emits batches matching \self.schema\ (with and without
projection).
- Retained end-to-end SQL aggregation regressions in \
ested_nullability.rs\ covering standard, distinct, and spilling aggregations.
##########
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:
Hi @alamb,
Thank you for pointing us to the PR review test coverage guide!
We refactored the test suite in \datafusion/common/src/nested_struct.rs\ and
\datafusion/core/tests/sql/aggregates/nested_nullability.rs\:
- Extracted shared test field builders (\ est_two_field_union\, etc.) to
eliminate repeated boilerplate across test cases.
- Streamlined unit tests to focus on distinct semantic cases: Sparse
adaptation, Dense adaptation, non-contiguous/reordered type-ID mappings, nested
Structs, mode mismatches, and field-set mismatches.
- Retained exact row-level unpacked value assertions (\10\, \\b\\, \30\),
active type IDs, and dense offsets without test repetition.
- Removed oversized duplicate tests from \
ested_nullability.rs\ to keep integration tests concise and focused on the
end-to-end bug report reproducer.
--
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]
