andygrove opened a new pull request, #5138:
URL: https://github.com/apache/datafusion-comet/pull/5138
## Which issue does this PR close?
Closes #5137.
## Rationale for this change
Arrow treats nested field nullability as part of `DataType` identity, so
`RecordBatch::try_new_with_options` rejects a column whose nested `nullable`
flags are narrower than the declared type — even though the data is fine, since
a non-null child is a strict subset of a nullable one. It is a metadata
assertion failure, not corruption.
Most Comet boundaries already normalize instead of asserting, which is why
this drift is usually invisible:
- `ScanExec::build_record_batch` casts every imported column to the declared
type, so the FFI boundary absorbs it.
- `SchemaAlignExec` casts the shuffle writer's input (added for #4515).
- `CometLocalTableScanExec` widens its output with `.asNullable` (added for
#4789).
Two stamp points still asserted, and because both sit inside native plans
(past `ScanExec`'s cast) nothing corrected the drift before them. The task
aborted deterministically for the partition's data shape, so retries did not
help:
```
Invalid argument error: column types must match schema types,
expected List(Struct("id": Int64, "flag": Boolean))
but found List(Struct("id": Int64, "flag": non-null Boolean))
at column index 0
```
- `ShuffleScanExec::poll_next` stamped the catalyst schema onto the decoded
shuffle block with no cast. This is the direct-read shuffle path, enabled by
default (`spark.comet.shuffle.directRead.enabled`).
- `ExpandStream::expand` derived its schema from `projections[0]` only and
stamped it onto the output of *every* projection, so it failed if any two
projections disagreed on a nested nullable flag for the same output column.
## What changes are included in this PR?
A new `cast_and_stamp_schema` helper in `datafusion-comet-common`, used at
both stamp points, which casts any column whose Arrow type differs from the
declared field type before stamping. Arrow's cast rewrites nested fields from
the target (`cast_struct_to_struct` builds with `to_fields`, `cast_list_values`
with `to.clone()`), which is exactly how `ScanExec` already absorbs this.
`ExpandExec::build_schema` now derives the output schema by widening nested
nullability across **all** projections instead of reading it off
`projections[0]`, so the declared schema is valid for every projection. The
runtime cast is a second, independent line of defence: each is tested on its
own, so neither silently carries the other.
Error messages now name the operator and the dotted path of the offending
column:
```
CometExpandExec cannot reconcile col[0] with its declared schema at
col_0.element.flag: expected nullable Boolean, found non-null Boolean
```
instead of only `at column index N`, which for a wide schema of deeply
nested structs was very hard to act on — the two printed types could differ by
one flag hundreds of characters in.
A narrowing cast (declared non-null child, actual data contains nulls) is
not silently accepted: `StructArray::try_new` rejects unmasked nulls for a
non-nullable field, so that case still errors loudly, now with the column path
named.
This PR is about the boundaries. The upstream nested-nullability drift
itself remains tracked by #4515, with concrete kernel instances in #4789 and
#4528.
## How are these changes tested?
New Rust unit tests, modelled on the existing `SchemaAlignExec` tests as the
issue suggested:
- `native/common/src/schema.rs` — 9 tests over `cast_and_stamp_schema`,
`describe_type_mismatch`, and `widen_nested_nullability`.
`stamps_nested_nullability_drift` first asserts that the plain stamp *does*
reject the drift, then that the helper absorbs it with values intact. That
assertion pins why the helper exists: if arrow ever relaxes nested nullability
comparison, it flips and the call sites can go back to a plain stamp.
- `native/core/src/execution/operators/expand.rs` — 5 tests: `build_schema`
widens across projections (driven from the narrow side, so it cannot pass by
echoing `projections[0]`), rejects projections of differing width, expands
divergent projections end-to-end, and reconciles even when handed the narrow
`projections[0]`-derived schema the planner used to build.
- `native/core/src/execution/operators/shuffle_scan.rs` — 2 tests: a decoded
block with a non-null nested child against a nullable declared type is
reconciled with values preserved, and an unreconcilable column names the
operator and column.
I confirmed these are genuine regression tests rather than tests that pass
either way: without the fix, `RecordBatch::try_new_with_options` rejects the
drift, which the first assertion in `stamps_nested_nullability_drift` now pins
permanently.
Also run and passing: full `cargo test --workspace` (135 + 33 native tests;
the 3 pre-existing `datafusion-comet-jni-bridge` `errors::tests` panic-handling
failures reproduce on a clean checkout of this branch point and are unrelated),
`cargo clippy --all-targets --workspace` clean, `cargo fmt`, and on Spark 3.5:
`CometExecSuite "expand operator"`, `CometNativeShuffleSuite` (27),
`CometShuffleSuite` (44), `CometAggregateSuite` (85, 2 canceled).
I did not reduce this to a Spark-level SQL test. Both sites need a native
plan where a kernel actually produces a nested type narrower than catalyst
declared, and those kernels are the subject of #4515/#4789/#4528 — so a SQL
repro would be pinned to whichever kernel currently drifts and would evaporate
when that kernel is fixed, which is the wrong thing to encode. The tests above
pin the boundaries directly, independent of which kernel drifts.
## Are these changes tested?
Yes.
## Are there any user-facing changes?
No API changes. Queries that previously aborted with `column types must
match schema types` in the direct-read shuffle path or in Expand (grouping sets
/ rollup / multiple distinct aggregates over nested columns) now complete.
Errors that remain genuinely unreconcilable are reported with the operator and
column path named.
## Follow-up
`ScanExec::build_record_batch` does the same cast-then-stamp but keeps its
own copy of the logic and its own `cast_time` metric, so its errors still
report only `at column index N`. Switching it to the shared helper would need
the helper to accept a timer; I left it out to keep this diff focused on the
two boundaries the issue names.
--
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]
