andygrove opened a new pull request, #6422: URL: https://github.com/apache/datafusion-comet/pull/6422
Backport of #5421 to `branch-1.0`. Cherry-picked from `6065705c16340c0be293212a71decfd9df4daae4`. The fix itself is unchanged. The Celeborn parts are dropped because `branch-1.0` does not have Celeborn shuffle planning; see "What changes are included" below. ## Which issue does this PR close? None on `branch-1.0`. #5421 is the `main` PR for #5419. Listed in #6201. ## Rationale for this change The bug ships in 1.0.0. A native partial aggregate can feed a Spark final aggregate that cannot read its buffer. `branch-1.0`'s `tagUnsafePartialAggregates` pass is the same as `main`'s before #5421: it decides before conversion whether the final aggregate will stay in Spark, and it skips the child-native check. So when the final falls back only because its input did not become native, the native partial below it survives. That happens when Comet shuffle is disabled, or when native shuffle is enabled but ineligible, for example because hash partitioning is disabled or the hash key is an array. - **AVG returns NULL.** A native AVG partial that sees no rows emits `(NULL, 0)`, and Spark's final merge turns the whole result into NULL. `branch-1.0` even marks non-decimal AVG as safe to mix. On the 1.0.0 release with `spark.comet.exec.shuffle.enabled=false`, `SUM(v), COUNT(v), AVG(v)` over a filter that leaves some scan partitions empty returns `[2.0, 1, null]`. That setting is what 1.0's own startup warning suggests when the Comet shuffle manager is not configured. - **`collect_list`, `collect_set` and `percentile` fail.** The native partial emits a buffer that Spark's final cannot read. The task fails with a `NullPointerException` in `UnsafeArrayData.foreach` for `collect_list` / `collect_set`, and with an `EOFException` for `percentile`. ## What changes are included in this PR? The fix is the original one; see #5421 for the details: - The single "mixed execution" opt-in becomes two independent ones. `supportsSparkPartialToNativeFinal` keeps the old meaning, for a Comet final or partial-merge that consumes a Spark buffer. `supportsNativePartialToSparkFinal` is new, for a Spark final that consumes a Comet buffer. MIN, MAX, COUNT, non-decimal non-TRY SUM, the bitwise aggregates, the bloom filter aggregate and `approx_count_distinct` opt in to the new direction. AVG stays out until #5420 fixes its empty state, and decimal SUM stays out because native precision overflow is sticky. - A new `revertUnsafePartialAggregates` pass runs on the converted plan. For each Spark final that would consume an unsafe native buffer, it restores the feeding partial, and any native merge stages and shuffles on that path, to Spark and reconverts the final's subtree. Native work below the partial, such as the scan and filter, stays native. The restored partial is tagged, so AQE's stage-only replanning keeps it in Spark. If the path stops at an existing query stage, the pass logs one warning and records an explain reason instead. The adaptations: - `CometExecRule`: `branch-1.0` does not have the Celeborn shuffle planning from #5537, so `preserveSparkAggregateBuffers` is dropped and `restoreSparkPartial` is used only by the new pass. The early tagging pass keeps `branch-1.0`'s Final-only consumer check (#5537 extended it to PartialMerge consumers on `main`) and only switches to the renamed predicate. - `RevertNativeForTransitionHeavyStages`: the #5421 hunk edits `hasUnsafeMixedAggregateAtStageBoundary`, which #5537 added and `branch-1.0` does not have, so it is dropped. The rule is off by default (`spark.comet.exec.transitionRevert.enabled`). - `CometCelebornShufflePlanningSuite` is not on `branch-1.0`, so its hunk is dropped. - `CometAggregateSuite` and `CometExecRuleSuite`: the #5421 tests with the imports they need. The FIRST/LAST percentile test that shows up in the `CometAggregateSuite` conflict comes from #5041, which is not on `branch-1.0` and must not be. Plans change only where a Spark final aggregate sits above a native partial: AVG and decimal SUM partials now run in Spark there, and COUNT partials can now stay native. There are no config or API changes. ## How are these changes tested? The tests from #5421, run locally on `branch-1.0` with JDK 17: - Default Spark 4.1 profile: `CometAggregateSuite` (110 tests) and `CometExecRuleSuite` (30 tests) pass. - Spark 3.5 / Scala 2.12 profile: the same two suites pass, 137 tests. The 3 existing map grouping-key tests are canceled there, as they are without this change. - The bug is present on `branch-1.0`, and the tests catch it. With the main-code changes reverted and the new `CometAggregateSuite` tests kept, 14 of 18 fail, with AQE both on and off. Decimal `AVG` returns `[null]` instead of `[200.000000]`, `COUNT(*), AVG(v)` returns `[1,null]` instead of `[1,1.0]`, and the `collect_list` / `collect_set` / `percentile` cases fail with the exceptions above. The 4 that pass are the COUNT and SUM controls, which keep their native partial either way. - `CometTPCDSV1_4_PlanStabilitySuite` and `CometTPCDSV2_7_PlanStabilitySuite` pass, 129 tests, so no TPC-DS golden plan changes. The aggregate SQL file tests (30) and the `CometExpressionSuite` explain tests pass. - Spotless and scalastyle on the default profile, scalafix in CHECK mode on Spark 3.5, and the syntactic scalafix check pass. PR CI on `branch-1.0` already runs the Comet suites on every Spark profile and the Spark SQL 3.5 and 4.1 suites, so the original's `run-all-spark-profiles` and `run-spark-4.1-tests` labels are not copied. `run-spark-4.0-tests` is added because COUNT partials can now feed a Spark final, and the case for that being safe rests partly on Spark 4.0's count-bug decorrelation. -- 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]
