andygrove opened a new pull request, #4861:
URL: https://github.com/apache/datafusion-comet/pull/4861
## Which issue does this PR close?
<!-- No dedicated issue. Follow-on to the mixed-execution buffer guard
(#1389). -->
N/A
This PR is stacked on top of #4750, which added the named-function fallback
reason. Please review and merge that PR first.
## Rationale for this change
When a Spark Final `HashAggregate` sits above a Spark (non-Comet) Partial
aggregate, Comet only runs the Final natively if every aggregate function's
intermediate buffer format is known to be compatible between Spark and Comet.
Until now only `MIN`, `MAX`, the bitwise aggregates, and `BloomFilter` were
marked compatible, so any plan containing `SUM` or `AVG` in this shape fell
back to Spark, for example:
```
HashAggregate [COMET: Spark Final aggregate without Comet Partial requires
compatible
intermediate buffer formats, but the following aggregate function(s) have
incompatible
buffers: sum]
```
This fallback appears across several TPC-DS plans (q10, q35, q45, q70/q70a).
The native buffers for these aggregates already match Spark's
`aggBufferAttributes`, so the fallback was overly conservative:
- Non-decimal `SUM` (Legacy/Ansi): single `sum` column, matching Spark.
- Decimal `SUM`: `(sum, is_empty)`, matching Spark's `(sum, isEmpty)` since
Spark always tracks `isEmpty` for decimal.
- Non-decimal `AVG`: `(sum: double, count: long)`, matching Spark.
The gate that blocked these was a coarse boolean that could not distinguish
data type or eval mode.
## What changes are included in this PR?
- Make the mixed-execution gate expression aware:
`CometAggregateExpressionSerde.supportsMixedPartialFinal` now takes the
aggregate function, and `QueryPlanSerde.supportsMixedExecution` passes it
through. The existing `MIN`/`MAX`/bitwise/`BloomFilter` overrides are updated
accordingly with no behavior change.
- Enable mixed partial/final execution for:
- `SUM` in any eval mode except `TRY` (covers decimal and non-decimal).
`TRY`-mode integer `SUM` is excluded because its native buffer carries a Comet
internal `has_all_nulls` column that Spark cannot read.
- Non-decimal `AVG`. Decimal `AVG` is deferred: its buffer layout matches,
but the decimal overflow path nulls the `count` column differently from Spark,
so it is left for a follow-up.
- `COUNT` remains excluded on purpose (unrelated to buffers: mixed `COUNT`
regressed AQE's `PropagateEmptyRelationAfterAQE` and the Spark 4.0 count-bug
decorrelation). The doc comment is corrected to reflect this.
- Regenerate the affected TPC-DS plan-stability golden files across all
Spark versions. The fallback reason is removed where `SUM`/`AVG` were the only
offenders, and shortened to `count` where `COUNT` is also present.
No native (Rust) changes are required.
## How are these changes tested?
- `CometExecRuleSuite`: new tests assert that `SUM` and non-decimal `AVG`
now allow mixed execution in both directions (Comet partial with Spark final,
and Spark partial with Comet final), mirroring the existing `MIN`/`MAX` tests,
plus a test asserting that `try_sum` remains blocked. The existing #1389
regression tests (which use `COUNT(*), SUM(id)`) still pass because `COUNT`
keeps them blocked; their comments are re-attributed to `COUNT`.
- `CometAggregateSuite`: end-to-end correctness tests force the aggregate
partial/final split across engines and confirm the results match Spark for
int/long/double `SUM` and `AVG` and for decimal `SUM`, in both split directions.
- Regenerated TPC-DS plan-stability golden files verify the resulting plans.
--
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]