comphead commented on PR #4791:
URL:
https://github.com/apache/datafusion-comet/pull/4791#issuecomment-5840913351
**Final aggregates fail on spill with `useLargeDataTypes=true`**
A job running this branch (rebased on `64e98918ab`, Spark 3.4) fails with:
```
org.apache.comet.CometNativeException: Invalid argument error: column types
must match schema types, expected Binary but found LargeBinary at column index 0
at org.apache.comet.Native.executePlan(Native Method)
...
at
org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleWriter.writeInternal(CometNativeShuffleWriter.scala:215)
```
Cause:
- `promote_byte_group_key` casts every Utf8/Binary group-by expression to
LargeUtf8/LargeBinary, in every aggregate mode.
- DataFusion 55.1's final hash aggregation builds its spill batches against
its *input* schema. In `datafusion-physical-plan`,
`AggregateHashTable<FinalMarker>::new` passes `agg.input().schema()` as the
state schema (`src/aggregates/aggregate_hash_table/final_table.rs:46`), and
`FinalSpillContext` gets the input stream's schema
(`src/aggregates/hash_stream.rs:1038-1044`). After a shuffle, that input still
declares `Binary`.
- When the final aggregate spills, `take_state_batch` stamps the
`LargeBinary` group values with that `Binary` field
(`src/aggregates/aggregate_hash_table/common.rs:320`). That is the error above.
Only tasks that spill fail.
- Partial aggregates are fine, because their state schema is derived from
the group-by expressions
(`src/aggregates/aggregate_hash_table/partial_table.rs:50-55`).
- The CUBE test doesn't catch this. It is `ignore`d and runs with an
unbounded pool, so no final aggregate spills.
Suggested fix: for DataFusion `Final` mode, cast the key columns in a
pass-through `ProjectionExec` below the aggregate, instead of casting the
group-by expressions. The aggregate's input, its group values and its spill
files then share one type. The aggregate expressions stay bound to the child
schema, because the state columns pass through unchanged. `Partial` and
`PartialMerge` keep `promote_byte_group_key`.
Two planner tests cover it:
- `final_aggregate_promotes_group_keys_at_its_input` checks that the input
to the final `AggregateExec` is the promoting projection and that its key type
matches the group-by output.
- `final_aggregate_with_large_group_keys_survives_spill` runs a final
aggregate over 50k binary keys with a 1 MiB pool, and asserts that it spills
and returns every key as `Binary`.
Both tests pass with the fix. With `Final` sent back through
`promote_byte_group_key`, the spill test fails with exactly the error above,
`ArrowError(InvalidArgumentError("column types must match schema types,
expected Binary but found LargeBinary at column index 0"))`, and the plan test
fails too.
Workaround until then:
`spark.comet.exec.aggregation.useLargeDataTypes=false`.
Two more things on this branch:
- The config doc says the default is false, but the code uses
`createWithDefault(true)`.
- When keys are promoted, the reverting `SchemaAlignExec` becomes the
aggregate's native root. It has no `metrics()`, so the HashAggregate's native
metrics, including spill counts, no longer reach Spark.
--
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]
