This is an automated email from the ASF dual-hosted git repository.
zhztheplayer pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 9b15c7de58 [GLUTEN-12668][CORE] Add injectPre hook so rules survive
whole-stage AQE fallback (#12669)
9b15c7de58 is described below
commit 9b15c7de586cac02f8d3f4d784b22db25d75a3b5
Author: BRIJ RAJ KISHORE <[email protected]>
AuthorDate: Tue Aug 4 14:35:02 2026 +0530
[GLUTEN-12668][CORE] Add injectPre hook so rules survive whole-stage AQE
fallback (#12669)
ExpandFallbackPolicy's whole-stage-fallback revert target (originalPlan in
HeuristicApplier.makeRule) is captured before any physical rule runs,
including injectPreTransform rules. A rule registered at injectPreTransform
therefore gets its rewrite stripped away whenever ExpandFallbackPolicy
promotes an individual-stage fallback to a whole-stage one, requiring a
second re-application at injectFinal as a workaround (see
RuntimeBloomFilterRewriteRule in #12151).
Adds a new injectPre hook to GlutenInjector.LegacyInjector, with its own
"pre" phase in HeuristicApplier.makeRule running before "transform".
fallbackPolicies now closes over the post-"pre" plan instead of the raw
originalPlan, so a whole-stage revert can no longer strip away anything
registered at injectPre.
Purely additive: a new empty-by-default builder list threaded through
HeuristicApplier's constructor. No existing backend needs to change unless
it opts in; backends-clickhouse builds and links unmodified against this
change. Once merged, #12151 will be rebased to register
RuntimeBloomFilterRewriteRule at injectPre only, collapsing its current two
registrations (injectPreTransform + injectFinal) down to one.
---
.../columnar/heuristic/HeuristicApplier.scala | 19 +++++++++++++++++--
.../gluten/extension/injector/GlutenInjector.scala | 6 ++++++
.../spark/sql/execution/FallbackStrategiesSuite.scala | 1 +
.../spark/sql/execution/FallbackStrategiesSuite.scala | 1 +
.../spark/sql/execution/FallbackStrategiesSuite.scala | 1 +
.../sql/gluten/GlutenFallbackStrategiesSuite.scala | 1 +
.../sql/gluten/GlutenFallbackStrategiesSuite.scala | 1 +
7 files changed, 28 insertions(+), 2 deletions(-)
diff --git
a/gluten-core/src/main/scala/org/apache/gluten/extension/columnar/heuristic/HeuristicApplier.scala
b/gluten-core/src/main/scala/org/apache/gluten/extension/columnar/heuristic/HeuristicApplier.scala
index cf1acaeed6..3c2d9ef998 100644
---
a/gluten-core/src/main/scala/org/apache/gluten/extension/columnar/heuristic/HeuristicApplier.scala
+++
b/gluten-core/src/main/scala/org/apache/gluten/extension/columnar/heuristic/HeuristicApplier.scala
@@ -32,6 +32,7 @@ import org.apache.spark.sql.execution.SparkPlan
*/
class HeuristicApplier(
session: SparkSession,
+ preBuilders: Seq[ColumnarRuleCall => Rule[SparkPlan]],
transformBuilders: Seq[ColumnarRuleCall => Rule[SparkPlan]],
fallbackPolicyBuilders: Seq[ColumnarRuleCall => SparkPlan =>
Rule[SparkPlan]],
postBuilders: Seq[ColumnarRuleCall => Rule[SparkPlan]],
@@ -47,10 +48,15 @@ class HeuristicApplier(
private def makeRule(call: ColumnarRuleCall): Rule[SparkPlan] = {
originalPlan =>
- val suggestedPlan = transformPlan("transform", transformRules(call),
originalPlan)
+ // `rewrittenPlan` (not `originalPlan`) is what `fallbackPolicies` uses
as its reference
+ // plan, so a whole-stage fallback reverts to the
already-`pre`-rewritten plan instead of
+ // the raw vanilla one. Rules registered here are therefore immune to
being reverted away
+ // by whole-stage fallback.
+ val rewrittenPlan = transformPlan("pre", preRules(call), originalPlan)
+ val suggestedPlan = transformPlan("transform", transformRules(call),
rewrittenPlan)
val finalPlan = transformPlan(
"fallback",
- fallbackPolicies(call).map(_(originalPlan)),
+ fallbackPolicies(call).map(_(rewrittenPlan)),
suggestedPlan) match {
case FallbackNode(fallbackPlan) =>
// we should use vanilla c2r rather than native c2r,
@@ -74,6 +80,15 @@ class HeuristicApplier(
new ColumnarRuleExecutor(phase, wrappedRules).execute(plan)
}
+ /**
+ * Rules applying to the plan before any offload decision is made. Unlike
`transformRules`,
+ * `preRules` are consistently baked into the plan that `fallbackPolicies`
reverts to on
+ * whole-stage fallback, so their effects cannot be reverted away.
+ */
+ private def preRules(call: ColumnarRuleCall): Seq[Rule[SparkPlan]] = {
+ preBuilders.map(b => b.apply(call))
+ }
+
/**
* Rules to let planner create a suggested Gluten plan being sent to
`fallbackPolicies` in which
* the plan will be breakdown and decided to be fallen back or not.
diff --git
a/gluten-core/src/main/scala/org/apache/gluten/extension/injector/GlutenInjector.scala
b/gluten-core/src/main/scala/org/apache/gluten/extension/injector/GlutenInjector.scala
index 012271485e..822607e187 100644
---
a/gluten-core/src/main/scala/org/apache/gluten/extension/injector/GlutenInjector.scala
+++
b/gluten-core/src/main/scala/org/apache/gluten/extension/injector/GlutenInjector.scala
@@ -46,6 +46,7 @@ class GlutenInjector private[injector] (control:
InjectorControl) {
object GlutenInjector {
class LegacyInjector {
+ private val preBuilders = mutable.Buffer.empty[ColumnarRuleCall =>
Rule[SparkPlan]]
private val preTransformBuilders = mutable.Buffer.empty[ColumnarRuleCall
=> Rule[SparkPlan]]
private val transformBuilders = mutable.Buffer.empty[ColumnarRuleCall =>
Rule[SparkPlan]]
private val postTransformBuilders = mutable.Buffer.empty[ColumnarRuleCall
=> Rule[SparkPlan]]
@@ -55,6 +56,10 @@ object GlutenInjector {
private val finalBuilders = mutable.Buffer.empty[ColumnarRuleCall =>
Rule[SparkPlan]]
private val ruleWrappers = mutable.Buffer.empty[Rule[SparkPlan] =>
Rule[SparkPlan]]
+ def injectPre(builder: ColumnarRuleCall => Rule[SparkPlan]): Unit = {
+ preBuilders += builder
+ }
+
def injectPreTransform(builder: ColumnarRuleCall => Rule[SparkPlan]): Unit
= {
preTransformBuilders += builder
}
@@ -86,6 +91,7 @@ object GlutenInjector {
private[injector] def createApplier(session: SparkSession):
ColumnarRuleApplier = {
new HeuristicApplier(
session,
+ preBuilders.toSeq,
(preTransformBuilders ++ Seq(
c => createHeuristicTransform(c)) ++ postTransformBuilders).toSeq,
fallbackPolicyBuilders.toSeq,
diff --git
a/gluten-ut/spark33/src/test/scala/org/apache/spark/sql/execution/FallbackStrategiesSuite.scala
b/gluten-ut/spark33/src/test/scala/org/apache/spark/sql/execution/FallbackStrategiesSuite.scala
index 07e4536689..91766cf1eb 100644
---
a/gluten-ut/spark33/src/test/scala/org/apache/spark/sql/execution/FallbackStrategiesSuite.scala
+++
b/gluten-ut/spark33/src/test/scala/org/apache/spark/sql/execution/FallbackStrategiesSuite.scala
@@ -185,6 +185,7 @@ private object FallbackStrategiesSuite {
transformBuilders: Seq[ColumnarRuleCall => Rule[SparkPlan]]):
HeuristicApplier = {
new HeuristicApplier(
spark,
+ Nil,
transformBuilders,
List(c => p => ExpandFallbackPolicy(c.caller.isAqe(), p)),
List(
diff --git
a/gluten-ut/spark34/src/test/scala/org/apache/spark/sql/execution/FallbackStrategiesSuite.scala
b/gluten-ut/spark34/src/test/scala/org/apache/spark/sql/execution/FallbackStrategiesSuite.scala
index 07e4536689..91766cf1eb 100644
---
a/gluten-ut/spark34/src/test/scala/org/apache/spark/sql/execution/FallbackStrategiesSuite.scala
+++
b/gluten-ut/spark34/src/test/scala/org/apache/spark/sql/execution/FallbackStrategiesSuite.scala
@@ -185,6 +185,7 @@ private object FallbackStrategiesSuite {
transformBuilders: Seq[ColumnarRuleCall => Rule[SparkPlan]]):
HeuristicApplier = {
new HeuristicApplier(
spark,
+ Nil,
transformBuilders,
List(c => p => ExpandFallbackPolicy(c.caller.isAqe(), p)),
List(
diff --git
a/gluten-ut/spark35/src/test/scala/org/apache/spark/sql/execution/FallbackStrategiesSuite.scala
b/gluten-ut/spark35/src/test/scala/org/apache/spark/sql/execution/FallbackStrategiesSuite.scala
index deebe20c4a..7616a62eed 100644
---
a/gluten-ut/spark35/src/test/scala/org/apache/spark/sql/execution/FallbackStrategiesSuite.scala
+++
b/gluten-ut/spark35/src/test/scala/org/apache/spark/sql/execution/FallbackStrategiesSuite.scala
@@ -186,6 +186,7 @@ private object FallbackStrategiesSuite {
transformBuilders: Seq[ColumnarRuleCall => Rule[SparkPlan]]):
HeuristicApplier = {
new HeuristicApplier(
spark,
+ Nil,
transformBuilders,
List(c => p => ExpandFallbackPolicy(c.caller.isAqe(), p)),
List(
diff --git
a/gluten-ut/spark40/src/test/scala/org/apache/spark/sql/gluten/GlutenFallbackStrategiesSuite.scala
b/gluten-ut/spark40/src/test/scala/org/apache/spark/sql/gluten/GlutenFallbackStrategiesSuite.scala
index f1227d3084..d7dd91a540 100644
---
a/gluten-ut/spark40/src/test/scala/org/apache/spark/sql/gluten/GlutenFallbackStrategiesSuite.scala
+++
b/gluten-ut/spark40/src/test/scala/org/apache/spark/sql/gluten/GlutenFallbackStrategiesSuite.scala
@@ -187,6 +187,7 @@ private object GlutenFallbackStrategiesSuite {
transformBuilders: Seq[ColumnarRuleCall => Rule[SparkPlan]]):
HeuristicApplier = {
new HeuristicApplier(
spark,
+ Nil,
transformBuilders,
List(c => p => ExpandFallbackPolicy(c.caller.isAqe(), p)),
List(
diff --git
a/gluten-ut/spark41/src/test/scala/org/apache/spark/sql/gluten/GlutenFallbackStrategiesSuite.scala
b/gluten-ut/spark41/src/test/scala/org/apache/spark/sql/gluten/GlutenFallbackStrategiesSuite.scala
index f1227d3084..d7dd91a540 100644
---
a/gluten-ut/spark41/src/test/scala/org/apache/spark/sql/gluten/GlutenFallbackStrategiesSuite.scala
+++
b/gluten-ut/spark41/src/test/scala/org/apache/spark/sql/gluten/GlutenFallbackStrategiesSuite.scala
@@ -187,6 +187,7 @@ private object GlutenFallbackStrategiesSuite {
transformBuilders: Seq[ColumnarRuleCall => Rule[SparkPlan]]):
HeuristicApplier = {
new HeuristicApplier(
spark,
+ Nil,
transformBuilders,
List(c => p => ExpandFallbackPolicy(c.caller.isAqe(), p)),
List(
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
