[branch-55] fix: adapt input batches with stricter nested nullability to planned schema in aggregation (#24394) - #24699
Conversation
…schema in aggregation (apache#24394) <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes apache#123` indicates that this PR will close issue apache#123. --> Closes apache#24069. <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. --> In DataFusion, in-memory table sources such as `MemTable::try_new` accept `RecordBatch`es whose schemas are stricter than the table's declared schema via `Schema::contains(&batches_schema)` (e.g. nullable nested fields declared on the table vs non-nullable nested fields in the input batches). However, `MemoryStream` previously advertised the declared table schema while emitting the underlying stricter `RecordBatch`es without adapting them. When downstream operators (such as `AggregateExec` with `array_agg` or distinct aggregation) received batches with stricter nested schemas, runtime type mismatch errors occurred (e.g. apache#24069). <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> 1. **`datafusion_common::nested_struct::adapt_batch_to_schema`**: - Adapts `RecordBatch`es whose nested schemas are stricter than a target schema. - Recursively reconstructs compatible nested Struct/List types. - Explicitly handles Dense and Sparse Union schema conformance while preserving type IDs and dense offsets without copying buffer data. - Requires exact Union type-ID sets and matching modes. - Does not broaden generic SQL CAST semantics (`requires_nested_struct_cast` remains untouched). 2. **`MemoryStream` Producer Boundary Normalization (`datafusion-physical-plan/src/memory.rs`)**: - Fixes the producer-side invariant exposed by apache#24069. - When batches have stricter schemas accepted by `MemTable::try_new`, `MemoryStream::poll_next` normalizes emitted batches using `adapt_batch_to_schema(batch, &self.schema)` whenever the runtime batch schema differs from `self.schema` and `self.schema.contains(batch.schema())`. - Ensures all `RecordBatch`es emitted by `MemoryStream` conform exactly to `stream.schema()`. 3. **Regression Coverage**: - Direct `MemoryStream` regressions in `memory.rs` verifying emitted batches match the advertised schema, including projection handling. - Unit tests in `nested_struct.rs` covering nested Struct and Dense/Sparse Union adaptation, unpacked scalar values, type IDs, offsets, reordered IDs, and incompatible Union layouts. - End-to-end SQL aggregation integration tests in `nested_nullability.rs` covering standard, DISTINCT, and spilling aggregations for apache#24069. <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? --> Yes: - `datafusion-common` unit tests for `adapt_batch_to_schema` and Union adaptation (`test_adapt_batch_to_schema_*`). - `datafusion-physical-plan` unit tests for `MemoryStream` emitted batch schema conformance and projection (`test_memory_stream_emitted_batch_matches_declared_schema*`). - `datafusion` core SQL integration tests in `datafusion/core/tests/sql/aggregates/nested_nullability.rs`. <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. --> No. Queries aggregating in-memory tables whose batches have stricter nested nullability than the table schema now succeed as expected. (cherry picked from commit 2326917)
github-actions
Bot
added
core
|
FYI @patrickswedish as the original author, would you mind reviewing this back port? |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## branch-55 #24699 +/- ##
=============================================
+ Coverage 81.16% 81.19% +0.02%
=============================================
Files 1110 1110
Lines 386911 387713 +802
Branches 386911 387713 +802
=============================================
+ Hits 314050 314788 +738
- Misses 54370 54405 +35
- Partials 18491 18520 +29 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Thanks for asking me to review the backport. I compared it against the merged #24394 and the actual branch-55 base rather than treating it as a mechanical cherry-pick. The Struct/Union path looks correctly carried over, and moving normalization to the MemoryStream producer boundary is still the right invariant. One backport-specific gap stood out: #24394 originally merged on top of #23914, which had already added recursive Map schema adaptation in nested_struct. branch-55 does not have that prerequisite. In this backport, adapt_batch_to_schema can now run for contained runtime-schema mismatches, but branch-55's cast_column still has no DataType::Map adaptation arm. So a stricter nested Map schema accepted by Schema::contains may fall through to Arrow's generic cast path rather than the recursive schema-conformance path present under the original PR. Before approval, I think it would be worth adding one focused regression: declared schema: Map<..., Struct> If Arrow 56 already handles this case correctly, then no extra code is needed. If it fails, we should either backport the minimal Map prerequisite from #23914 or explicitly constrain the branch-55 adapter to shapes this branch can safely normalize. Everything else I checked looks aligned with the original fix direction. |
alamb
left a comment
alamb
left a comment
There was a problem hiding this comment.
Thanks @timsaucer and @patrickswedish
maybe @kosiew the original reviewer can weigh in here too
kosiew
left a comment
kosiew
left a comment
There was a problem hiding this comment.
Thanks for working on this backport. The overall approach looks good to me. I left a couple of non-blocking suggestions for additional regression coverage around Union and Map schema adaptation.
| @@ -673,4 +684,128 @@ mod lazy_memory_tests { | |||
|
|
|||
There was a problem hiding this comment.
The new Union adaptation has good coverage at the helper level, but I think it would also be useful to exercise it through MemoryStream, where we enforce the producer-side schema invariant. Could we add a small Union case with a nullable declared child and a non-nullable runtime child, then assert that emitted_batch.schema() == stream.schema()? That would give us coverage for the Union reconstruction path at the actual integration boundary.
| target_mode, | ||
| cast_options, | ||
| ), | ||
| _ => Ok(cast_with_options(source_col, target_type, cast_options)?), |
There was a problem hiding this comment.
Could we add a focused MemoryStream regression for a contained Map<..., Struct> where the runtime nested field is non-nullable and the declared nested field is nullable? My understanding is that this should already work because shapes accepted by Schema::contains retain the same Map entry/key/value structure, and Arrow's generic Map cast should recursively cast the value Struct and rebuild it with the target fields. A regression here would confirm that assumption for this branch.
If it fails, we should either backport the minimal Map prerequisite from
https://github.com/apache/datafusion/pull/23914 or explicitly constrain
the branch-55 adapter to shapes this branch can safely normalize.
as suggested by @patrickswedish
…tation Adds two `MemoryStream` regression tests requested during review of the apache#24394 backport: - Dense Union with nullable declared children and non-nullable runtime children, exercising the Union reconstruction path at the producer boundary and asserting preserved type IDs and child values. - `Map<Utf8, Struct<v>>` where the runtime nested field is non-nullable and the declared nested field is nullable. This passes on branch-55 without a `DataType::Map` arm in `cast_column`: `Schema::contains` accepts the stricter shape and Arrow's generic Map cast recursively casts the value Struct and rebuilds it with the target fields. No backport of the Map prerequisite from apache#23914 is required. Both assert `emitted_batch.schema() == stream.schema()`. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
5 participants
This is a back port of #24394 for
branch-55.The original issue is #24069
Note that the original PR was marked as an auto-detected api change, but I believe that is a false positive.