andygrove opened a new pull request, #5159:
URL: https://github.com/apache/datafusion-comet/pull/5159
## Which issue does this PR close?
Closes #5158.
## Rationale for this change
`SparkCollectList` / `SparkCollectSet` derive their `return_type` and
`state_fields` from the aggregate argument's *declared* type, but the
accumulator builds its list from the arrays it actually received —
`ScalarValue::new_list` honors its `data_type` parameter only when the value
set is empty, and otherwise just concatenates the collected scalars. Nothing
forces the two to agree, so an argument expression whose `evaluate()` output
type differs from its `data_type()` in nested field nullability makes the
grouped `AggregateExec` fail validating its own output batch:
```
Invalid argument error: column types must match schema types,
expected List(Struct("flag": Boolean, "items": List(Struct("n": Int32))))
but found List(Struct("flag": non-null Boolean, "items": List(Struct("n":
non-null Int32))))
at column index 1
```
`coerce_collect_child_nullability` (added in #4720) already casts the
argument to the all-nullable variant of its type, but returns early when the
declared type is *already* all-nullable — which is exactly the case above,
where the declared type is all-nullable and the runtime array is not. The cast
only helped in the opposite direction.
## What changes are included in this PR?
Cast unconditionally for nested argument types, so the accumulator is
guaranteed to see arrays of exactly the type the plan declared, whichever way
the drift goes.
Two things worth noting for review:
* **Not observable from Spark.** The declared output type is unchanged in
every case. Where the argument type was already all-nullable,
`make_all_fields_nullable` is the identity, so the aggregate's declared type is
byte-identical to before and only the runtime array is normalized. Where it was
not, the cast was already being inserted.
* **Free when there is no drift.** `cast_with_options` returns a shallow
clone when the source and target types are equal, so the added cast costs
nothing in the common case. When they do differ it is still metadata-only:
`cast_struct_to_struct` and `cast_list_values` stamp the target fields and
clone the child buffers rather than copying data.
The rationale comment is also refreshed. It described the accumulator as
marking all element fields nullable regardless of the input; that is no longer
true in DataFusion 54.1.0 — the accumulator preserves the input array's type,
so declared and produced already agree when there is no drift, and the coercion
was a normalization rather than a defence.
## How are these changes tested?
Three new unit tests in `planner.rs`:
* `test_collect_agg_absorbs_nested_nullability_drift` — builds `collect_set`
and `collect_list` against a schema with nullable nested leaves and feeds them
a batch whose array has non-nullable leaves, then compares
`state_fields()[0].data_type()` against the type of the emitted state array.
That is exactly the pair `RecordBatch::try_new` compares inside
`GroupedHashAggregateStream::emit`. Asserts the drift is present without the
coercion and absent with it. Without the fix this test fails, and reproduces
the error message above verbatim.
* `test_collect_agg_coercion_widens_all_nested_fields` — the declared
element type is the all-nullable variant of the argument's declared type.
* `test_collect_agg_no_coercion_for_primitive_child` — primitives carry no
field-level nullability, so no cast is inserted.
Existing coverage: `cargo test -p datafusion-comet --lib` (135 passed) and
`CometAggregateSuite` under `-Pspark-3.5` (86 passed, 0 failed, 2 canceled as
Spark-4.1-gated), including the existing `collect_list over struct with
non-nullable fields` test which covers the original direction.
## Follow-ups (not in this PR)
* The root cause on the Comet side is that `CreateNamedStruct` recomputes
its `Fields` from `batch.schema()` in `evaluate()` but from the plan-time
schema in `data_type()`. Resolving the `Fields` once at construction would make
the expression incapable of drifting, and drop per-batch field construction.
That is a signature change with wider blast radius, so it is better separate.
* The underlying defect is upstream: `ScalarValue::new_list` accepts a
`data_type` it then discards for non-empty input, so no accumulator built on it
can guarantee the type its `state_fields()` advertises. Filed against
DataFusion; if fixed there, this coercion can be dropped entirely.
--
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]
