feat: Stage based fallback by karuppayya · Pull Request #4519 · apache/datafusion-comet

karuppayya · 2026-05-29T23:47:08Z

Which issue does this PR close?

Closes #4518 .

Rationale for this change

Introduces a stage-aware fallback that reverts stages to Spark row-based execution when the number of C2R transitions exceeds a threshold.

What changes are included in this PR?

RevertNativeForTransitionHeavyStages — postColumnarTransitions rule that counts C2R transitions per stage and reverts via CometExec.originalPlan when threshold exceeded
Configs: spark.comet.exec.transitionRevert.enabled (default: true), spark.comet.exec.transitionRevert.maxTransitions (default: 2, min: 2)
Works for both AQE and non-AQE paths

How are these changes tested?

Unit tests added

karuppayya · 2026-06-01T15:30:31Z

cc: @mbutrovich @andygrove @mbutrovich @parthchandra
The main idea here is for the stage to fallback entirely to JVM based on cost(number of transition between Columnar and row).

andygrove · 2026-06-01T18:15:16Z

Thanks @karuppayya. I've observed this issue as well. I introduced CometConf.COMET_EXEC_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED to control whether a row-based stage gets converted to columnar shuffle. It currently defaults to true, but I am curious whether changing the default to false would help with this. I was planning on running benchmarks with this config disabled but have not done that yet.

andygrove · 2026-06-01T18:20:04Z

@karuppayya new test suites have to be manually added to pr_build_linux.yml and pr_build_macos.yml in .github workflows.

Suite not found in workflow .github/workflows/pr_build_linux.yml: org.apache.comet.rules.RevertNativeForTransitionHeavyStagesSuite

karuppayya · 2026-06-01T19:10:12Z

Thanks @andygrove . I think difference is that COMET_EXEC_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED only checks the immediate child of the shuffle to decide whether to convert or not. It doesn't look at the full stage plan, so you could still end up with costly transitions deeper in the stage

This changes tries evaluates the full stage plan and makes a per-stage decision based on the actual transition cost. In future we can extend it to factor in estimated row counts, operator complexity, etc.

karuppayya · 2026-06-01T19:11:32Z

Also just curious since i dont have the background, why is the CI trigger manual?

comphead · 2026-06-01T19:40:24Z

Also just curious since i dont have the background, why is the CI trigger manual?

Comet has a large ASF consumption, and we working to offload longrunning CI flows to the user forks. For now we trying to run less pipelines in parallel

comphead

Thanks @karuppayya

One thing to mention is: Comet mostly works on planning phase to inject Comet nodes to be executed. AQE, DPP also.

However this approach proposes a totally new vision and make decisions in runtime? Spark stage is the runtime unit and might also depend on data distribution.

This is a lot of gray area here, for example if Spark scans a folder with 100 files which are unevenly distributed, leading spark tasks to be executed with different runtime and workload, possibly leading to overloaded single task to work more than others. Is it considered a heavy weight stage? what is the strategy to fallback? restart the task?

I just want to understand the approach better and see whats pros and cons it could bring up

karuppayya · 2026-06-02T00:19:18Z

Thanks @comphead.

However this approach proposes a totally new vision and make decisions in runtime?

Yes and no. This is still a planning-phase rule (postColumnarTransitions). For non-AQE, it's staticsince the decision is made once during initial plan optimization. For AQE, it runs per-stage as each stage is planned at every shuffle boundary during runtime, but it's still a plan based decision.

Is it considered a heavy weight stage? what is the strategy to fallback?

Currently the heuristic is structural ie it is based on counts the number of C2R transitions in a stage's plan. It doesn't take into account data volume etc. Incorporating those kinds of statistics could be a natural next step toward a more cost-based decision.

comphead · 2026-06-02T16:44:59Z

Thanks @karuppayya for the explanation, having conversion counter guardrail on planner makes a lot of sense to me

parthchandra · 2026-06-02T23:30:59Z

+        _ >= 2,
+        "Must be >= 2. A reverted stage still requires at least one " +
+          "R2C at the columnar shuffle boundary.")
+      .createWithDefault(2)


This is probably too aggressive and is likely to cause CI failure (because the plans will not match).

I'd like to keep the threshold at 2 eventually since more than 2 transitions means significant overhead. But for now I've bumped it to 5 for CI to pass. Once we've reviewed the logic , we can iterate on the default and fix test failures than came with it?.

parthchandra · 2026-06-02T23:34:03Z

+      case RowToColumnarExec(child) => child
+    }
+    val reverted = stripped.transformUp {
+      case cometShuffle: CometShuffleExchangeExec =>


CometBroadcastExchange extends CometPlan (not CometExec) and will not be matched by any case in this pass.

In AQE this is already handled — the broadcast exchange is wrapped in BroadcastQueryStageExec (which extends QueryStageExec), and countTransitions stops at QueryStageExec.
For non-AQE, I've added BroadcastExchangeLike to the stop condition in countTransitions to treat it as an execution boundary.(as broadcast runs as a separate job independentof the parent stage).

parthchandra · 2026-06-02T23:42:58Z

+      case cometExec: CometExec =>
+        if (cometExec.originalPlan.children.size == cometExec.children.size) {
+          cometExec.originalPlan.withNewChildren(cometExec.children)
+        } else {


When does this get hit? We should probably never reach this in which case a warning in the log is called for.

Added a warnign

parthchandra · 2026-06-02T23:46:22Z

+    plan.transformUp { case exchange: ShuffleExchangeLike =>
+      revertStageIfNeeded(exchange.child, exchange.supportsColumnar)
+        .map(reverted => exchange.withNewChildren(Seq(reverted)))
+        .getOrElse(exchange)


Why is there no case _ => arm like in the AQE case.

parthchandra · 2026-06-02T23:49:42Z

+      }
+    }
+  }
+}


Can we add a test that reverts and then executes the plan? Maybe two variants like - fallback triggered by unsupported expressions that are reverted, and a threshold set to 1 to force revert.
Test should verify results against Spark.

I disabled the project operator to guarantee transitions Do you think we still need a test with a naturally fallback (unsupported expression may be UDF), or is the current approach
sufficient?

parthchandra · 2026-06-02T23:55:01Z

+    plan.transformUp {
+      case node if !node.isInstanceOf[QueryStageExec] && !node.supportsColumnar =>
+        val newChildren = node.children.map { child =>
+          if (child.supportsColumnar) ColumnarToRowExec(child) else child


You could also hit a Row child and a Columnar parent (the reverse transition).

After revertToSpark, all CometExec nodes are replaced with their row-based operator. The only columnar nodes remaining would be QueryStageExec leaves (inputs from prior stages). I think we will have only columnar children feeding row-based parents. May be i am missing the scenario, an example would help.

parthchandra · 2026-06-02T23:59:11Z

+        if (newChildren != node.children) node.withNewChildren(newChildren) else node
+    }
+  }
+}


You might want to remove CometTelemetryExec as well. It won't do any harm if it is there but it's pointless to collet Comet telemetry for Spark nodes.

karuppayya

Thanks @parthchandra for the rview.
i have also cleaned up tets.

karuppayya · 2026-06-03T17:28:44Z

+      case RowToColumnarExec(child) => child
+    }
+    val reverted = stripped.transformUp {
+      case cometShuffle: CometShuffleExchangeExec =>


In AQE this is already handled — the broadcast exchange is wrapped in BroadcastQueryStageExec (which extends QueryStageExec), and countTransitions stops at QueryStageExec.
For non-AQE, I've added BroadcastExchangeLike to the stop condition in countTransitions to treat it as an execution boundary.(as broadcast runs as a separate job independentof the parent stage).

karuppayya · 2026-06-03T18:01:37Z

+    plan.transformUp { case exchange: ShuffleExchangeLike =>
+      revertStageIfNeeded(exchange.child, exchange.supportsColumnar)
+        .map(reverted => exchange.withNewChildren(Seq(reverted)))
+        .getOrElse(exchange)


karuppayya · 2026-06-03T18:02:06Z

+      case cometExec: CometExec =>
+        if (cometExec.originalPlan.children.size == cometExec.children.size) {
+          cometExec.originalPlan.withNewChildren(cometExec.children)
+        } else {


Added a warnign

karuppayya · 2026-06-03T18:38:02Z

+    plan.transformUp {
+      case node if !node.isInstanceOf[QueryStageExec] && !node.supportsColumnar =>
+        val newChildren = node.children.map { child =>
+          if (child.supportsColumnar) ColumnarToRowExec(child) else child


After revertToSpark, all CometExec nodes are replaced with their row-based operator. The only columnar nodes remaining would be QueryStageExec leaves (inputs from prior stages). I think we will have only columnar children feeding row-based parents. May be i am missing the scenario, an example would help.

karuppayya · 2026-06-03T18:40:53Z

+        if (newChildren != node.children) node.withNewChildren(newChildren) else node
+    }
+  }
+}


karuppayya · 2026-06-03T18:46:47Z

+        _ >= 2,
+        "Must be >= 2. A reverted stage still requires at least one " +
+          "R2C at the columnar shuffle boundary.")
+      .createWithDefault(2)


I'd like to keep the threshold at 2 eventually since more than 2 transitions means significant overhead. But for now I've bumped it to 5 for CI to pass. Once we've reviewed the logic , we can iterate on the default and fix test failures than came with it?.

karuppayya · 2026-06-03T19:57:35Z

+      }
+    }
+  }
+}


I disabled the project operator to guarantee transitions Do you think we still need a test with a naturally fallback (unsupported expression may be UDF), or is the current approach
sufficient?

comphead · 2026-06-04T01:32:58Z

+      stagePlan: SparkPlan,
+      outputColumnar: Boolean): Option[SparkPlan] = {
+    val transitionCount = countTransitions(stagePlan)
+    if (transitionCount <= maxTransitions) return None


To make a user friendly response it would be better to use withFallbackReason instead of logging. Example CometExecRule.scala

karuppayya force-pushed the COMET-4518 branch 2 times, most recently from 9061d20 to 7b69f7c Compare June 1, 2026 19:06

comphead reviewed Jun 1, 2026

View reviewed changes

karuppayya added 3 commits June 2, 2026 11:49

Stage based fallback

0f6db90

Add suites to yml

3ed99e4

Fix tests

3b51521

karuppayya force-pushed the COMET-4518 branch from caac4c0 to 3b51521 Compare June 2, 2026 18:49

parthchandra reviewed Jun 2, 2026

View reviewed changes

Address review comments

a0a039f

karuppayya commented Jun 3, 2026

View reviewed changes

comphead reviewed Jun 4, 2026

View reviewed changes

Conversation

karuppayya commented May 29, 2026

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

How are these changes tested?

Uh oh!

karuppayya commented Jun 1, 2026

Uh oh!

andygrove commented Jun 1, 2026

Uh oh!

andygrove commented Jun 1, 2026

Uh oh!

karuppayya commented Jun 1, 2026

Uh oh!

karuppayya commented Jun 1, 2026

Uh oh!

comphead commented Jun 1, 2026

Uh oh!

comphead left a comment

Choose a reason for hiding this comment

Uh oh!

karuppayya commented Jun 2, 2026

Uh oh!

comphead commented Jun 2, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

karuppayya left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

comphead Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

comphead Jun 4, 2026 •

edited

Loading