andygrove opened a new pull request, #5213:
URL: https://github.com/apache/datafusion-comet/pull/5213
## Which issue does this PR close?
Part of #5199 (item 5: `EliminateRedundantTransitions.hasCometNativeChild`
is quadratic).
## Rationale for this change
`EliminateRedundantTransitions` asks `hasCometNativeChild` whether the
subtree under a
`ColumnarToRowExec` contains a Comet operator, and that check is a
`TreeNode.exists` scan of the
subtree. Because the rule runs as a `transformUp`, the scan is repeated at
every
`ColumnarToRowExec` and each one re-walks everything below it, so stacked
transitions rescan the
same nodes and the combined cost is quadratic in plan size. The rule runs on
the driver for every
query, and again for every query stage under AQE.
The scan stops at the first Comet operator it finds, so the quadratic
behavior only shows up over
a Comet-free subtree. That case is real: the rule runs on every plan,
including one where Comet
took nothing but Spark still inserted columnar transitions.
Measured with a throwaway probe that applies the rule to a chain of
`ColumnarToRowExec` /
`RowToColumnarExec` pairs over a vanilla (non-Comet) Parquet scan, counting
the nodes visited by
the subtree scan and timing the rule (mean of 20 applications after warm-up,
Spark 4.1 / JDK 17):
| plan nodes | before: nodes visited | before | after: nodes visited | after
|
| ---------- | --------------------- | ------ | -------------------- | -----
|
| 203 | 10,302 | 0.30 ms | 202 | 0.07
ms |
| 803 | 161,202 | 4.62 ms | 802 | 0.13
ms |
| 1,603 | 642,402 | 17.20 ms | 1,602 | 0.23
ms |
| 3,203 | 2,564,802 | 66.52 ms | 3,202 | 0.43
ms |
| 6,403 | 10,249,602 | 264.17 ms | 6,402 | 0.86
ms |
When the subtree does contain a Comet operator both versions are already
linear and measure the
same: the rule rewrites the innermost `ColumnarToRowExec` into a
`CometColumnarToRowExec`, which
is itself a `CometPlan`, so every ancestor's scan short-circuits a node or
two down.
## What changes are included in this PR?
Memoizes the subtree scan for the duration of a single rule invocation.
- The memo is keyed on identity, because `SparkPlan` equality and hash are
themselves subtree
walks and a `HashMap` would reintroduce the quadratic cost in the lookup.
- It is a local of `_apply` rather than a field, because the rule instance
lives for the whole
session and must not retain plans.
- `transformUp` preserves the identity of subtrees it does not rewrite, so a
rebuilt node still
hits the memo one level down.
Behavior is unchanged. `containsCometPlan` reproduces
`op.exists(_.isInstanceOf[CometPlan])`
exactly, and `QueryStageExec` / `ReusedExchangeExec` are still unwrapped
only at the root of the
checked subtree, as before.
## How are these changes tested?
Existing coverage: every Comet query test exercises this rule, and
`RevertNativeForTransitionHeavyStagesSuite` / `CometMapInBatchSuite` assert
on its output
directly.
Added `EliminateRedundantTransitionsSuite`, which puts a Comet columnar
branch and a Spark-only
columnar branch under one plan and asserts the rule rewrites only the Comet
one. A memo that
leaked a result between branches would fail it. The probe used for the
numbers above was
throwaway and is not included.
--
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]
