andygrove opened a new pull request, #6609: URL: https://github.com/apache/datafusion-comet/pull/6609
## Which issue does this PR close? Closes #2785. Related: #5207 (whole-stage revert by default). Earlier attempts at this: #2906, #3220 and #3816. ## Rationale for this change When Comet accelerates part of a query stage and then has to hand every row to Spark, the conversion from columnar batches to rows can cost more than the native operators save. The clearest case is a native scan over a very large table followed immediately by a fallback, where the rest of the stage is a trivial projection: Comet converts every scanned row and gains almost nothing in return. Comet already has a per-stage revert (`spark.comet.exec.transitionRevert.enabled`), but it decides by counting transitions. A count cannot tell a transition that converts billions of scanned rows from one that converts a few aggregated rows, which is why #5207 notes that reverting on any transition also throws away native scans, filters and aggregates that are worth keeping. This PR adds an experimental cost model that makes the same per-stage decision by estimated cost instead: it compares the stage as Comet would run it with the same stage reverted to Spark, weighting each operator and transition by the rows it is estimated to process. ## What changes are included in this PR? Everything is off by default (`spark.comet.exec.costModel.enabled=false`). **Pluggable cost model.** `org.apache.comet.cost.CometCostModel` has one method, `estimate(cometPlan, sparkPlan): CometCostEstimate`, which is given both candidate plans for one query stage and returns a cost for each. `spark.comet.exec.costModel.class` selects the implementation, so workload-specific models can be tried without changing Comet, as #2785 asks. **Default model.** `DefaultCometCostModel` charges each operator `input rows x weight`: - Spark operators have a per-row weight by kind (scan, filter, project, aggregate, join, sort, window). Scan and transition weights scale with schema width, and filter and project weights with the number of expressions they compute. - A Comet native operator costs the Spark operator it replaces divided by `spark.comet.exec.costModel.nativeSpeedup` (default 2.0). - A conversion to or from Comet's columnar batches costs `spark.comet.exec.costModel.transitionCostFactor` (default 3.0) times what Spark pays to convert its own vectorized scan output to rows. - Row counts come from the runtime statistics of materialized query stages under AQE (exact), and otherwise from the logical statistics of the stage's scans, pushed through the stage with fixed selectivities. **Decision.** `RevertNativeForTransitionHeavyStages` reverts a stage when `sparkCost / cometCost` is below `spark.comet.exec.costModel.minSpeedup` (default 1.0). It reuses the existing revert machinery, so the cost model inherits stage scoping under AQE and non-AQE, the guard that keeps an aggregate from being split between engines, and plan-only preview support. The fallback reason reports the estimated speedup and names the configs. A model that throws or cannot be loaded leaves the Comet plan in place. With the defaults: | Stage | Estimated speedup | Decision | | --- | --- | --- | | native scan -> C2R -> Spark project | 0.76 | revert | | native scan -> native filter -> C2R -> Spark project | 2.16 | keep | | native scan -> native partial aggregate -> shuffle | 2.56 | keep | | shuffle read -> native final aggregate -> C2R | 2.00 | keep | The filter case shows the property a transition count misses: Spark always pays its own columnar-to-row conversion on every scanned row, while Comet only pays at the point of fallback, after the native filter has removed most rows. **How this differs from the earlier attempts.** #2906 averaged a per-operator acceleration over the plan and ignored row counts. #3220 used one row count for the whole plan and could not revert scans. #3816 was an operator-count coverage threshold over the whole query. This PR decides per stage, weights by estimated rows per operator, compares two concrete plans (so the cost of bridging a reverted stage back to a Comet shuffle is included), and can revert native scans. **Fixes to the stage revert that the cost model depends on.** A cost model reverts stage shapes that the transition count rarely or never reached, and two of them produced plans that fail at runtime. Both are in code that `transitionRevert` shares, where they need unusual settings to hit (the first needs `maxTransitions=0`): - A stage that reverts to a bare vectorized Spark scan (for example `SELECT * FROM t`) was left without a columnar-to-row transition and failed with a `ClassCastException` (`ColumnarBatch` cannot be cast to `InternalRow`). - A reverted stage feeding a Comet native shuffle was wrapped in Spark's `RowToColumnarExec`, whose vectors the native shuffle writer cannot read (`OffHeapColumnVector` cannot be cast to `CometVector`). It is now bridged with `CometSparkToColumnarExec`, and the stage is not reverted if that conversion does not support the schema. These could be split into their own PR if preferred. **Docs.** New "Cost-Based Stage Fallback (Experimental)" section in the transitions tuning guide. ### Limitations and open questions - **The default weights are not calibrated.** They are initial estimates, chosen so that the motivating case reverts and stages with real native work do not. The next step is to calibrate `nativeSpeedup`, `transitionCostFactor` and the per-operator weights against TPC-H/TPC-DS and the existing columnar-to-row microbenchmarks (#5112 shows the conversion cost varies a lot with schema and batch size). Feedback on whether the knobs exposed as configs are the right ones is welcome. - **A fully native stage that outputs every scanned row is treated like the motivating case.** `SELECT a + 1 FROM t` with a native projection still converts every row for its consumer, so with the defaults it reverts too. That is consistent with the model, but whether it matches reality depends on the calibration above. - **Cardinality estimates are coarse.** Without AQE runtime statistics, row counts come from table size and fixed selectivities. Because every cost is linear in rows, a single-input stage only depends on row counts relative to its input, which makes the model less sensitive to this than it sounds. Spark CBO column statistics are not used yet. - **The stage is costed before `EliminateRedundantTransitions` runs.** A `RowToColumnarExec` that a Comet JVM shuffle reads through without converting is still charged, to both candidates equally. - The known issues in `revertToSpark` (#5719, #6152) apply to stages reverted by the cost model as they do to stages reverted by the transition count. ## How are these changes tested? New `CometCostModelSuite`, registered in the Linux and macOS workflows, covering with AQE on and off: - the motivating case reverts and returns correct results, with the expected fallback reason - a native filter or aggregate ahead of the fallback keeps the stage in Comet - the estimates the default model produces, and how `transitionCostFactor` and `nativeSpeedup` move them - `minSpeedup` in both directions - a custom model class decides which stages revert - a model that throws, or a class that cannot be loaded, leaves the Comet plan in place - the mixed-aggregate guard still applies when the cost model asks for a revert - scan-only stages revert correctly as a result stage and when feeding a Comet native or JVM shuffle (the two fixes above) `RevertNativeForTransitionHeavyStagesSuite` gains a test that reverts a scan-only stage through the transition count, and its existing tests pass unchanged. The change also compiles against Spark 3.4, and against Spark 3.5 with strict warnings and the scalafix check. The cost model is off by default and the rule returns early when neither it nor `transitionRevert` is enabled, so default plans are unchanged. No benchmarks have been run yet. -- 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]
