This is an automated email from the ASF dual-hosted git repository.
peter-toth pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 3ad5b64dbca1 [SPARK-57194][SQL] Add preOperatorOptimizationRules
extension point to Optimizer
3ad5b64dbca1 is described below
commit 3ad5b64dbca11d6ba37bfdf6904e53ee2c00fd6a
Author: Shrirang Mhalgi <[email protected]>
AuthorDate: Wed Jun 10 11:57:36 2026 +0200
[SPARK-57194][SQL] Add preOperatorOptimizationRules extension point to
Optimizer
### What changes were proposed in this pull request?
Add `preOperatorOptimizationRules` - a new optimizer extension point that
runs in a `Once` batch between "Aggregate" and the main fixed-point "Operator
Optimization" batch. Wired through
`SparkSessionExtensions.injectPreOperatorOptimizationRule` and
`BaseSessionStateBuilder`. The batch is a no-op when no rules are registered.
### Why are the changes needed?
Custom rules injected via `injectOptimizerRule` run inside the fixed-point
batch alongside built-in rules (FoldablePropagation, ConstantFolding,
PushDownPredicates). A rule that needs to observe the original plan shape
(e.g., cross-side join predicates before they are folded into single-side
constants) is silently defeated when a built-in rule transforms the plan first
within the same fixed-point iteration.
None of the existing extension points cover this:
- `extendedOperatorOptimizationRules` - same fixed-point batch
- `postHocResolutionRules` - analyzer phase, too early
- `earlyScanPushDownRules` - runs after optimization, scoped to scan
pushdown
- `preCBORules` - runs after operator optimization, too late
### Does this PR introduce _any_ user-facing change?
Yes. New public API will be available:
`SparkSessionExtensions.injectPreOperatorOptimizationRule(builder:
RuleBuilder)`.
### How was this patch tested?
- New test in `SparkSessionExtensionSuite` verifying the injected rule is
wired into the `preOperatorOptimizationRules` accessor.
- All existing SparkSessionExtensionSuite tests pass.
### Was this patch authored or co-authored using generative AI tooling?
Yes.
Closes #56277 from shrirangmhalgi/SPARK-57194-early-optimizer-rules.
Authored-by: Shrirang Mhalgi <[email protected]>
Signed-off-by: Peter Toth <[email protected]>
---
.../apache/spark/sql/catalyst/optimizer/Optimizer.scala | 12 ++++++++++++
.../org/apache/spark/sql/SparkSessionExtensions.scala | 17 +++++++++++++++++
.../spark/sql/internal/BaseSessionStateBuilder.scala | 13 +++++++++++++
.../apache/spark/sql/SparkSessionExtensionSuite.scala | 6 ++++++
4 files changed, 48 insertions(+)
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
index a1fe4cd8674b..42914e98aa6d 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
@@ -228,6 +228,10 @@ abstract class Optimizer(catalogManager: CatalogManager)
Batch("Aggregate", fixedPoint,
RemoveLiteralFromGroupExpressions,
RemoveRepetitionFromGroupExpressions),
+ // Injected rules run once here, before the operator-optimization fixed
point, so they
+ // can observe the plan before FoldablePropagation/ConstantFolding rewrite
it.
+ Batch("Pre Operator Optimization", Once,
+ preOperatorOptimizationRules: _*),
operatorOptimizationBatch,
Batch("Clean Up Temporary CTE Info", Once, CleanUpTempCTEInfo),
// This batch rewrites plans after the operator optimization and
@@ -507,6 +511,14 @@ abstract class Optimizer(catalogManager: CatalogManager)
*/
def extendedOperatorOptimizationRules: Seq[Rule[LogicalPlan]] = Nil
+ /**
+ * Override to provide additional rules that run in a single pass before the
main
+ * operator optimization fixed-point batch. Use this for rules that need to
observe the plan
+ * as it enters the operator-optimization fixed point (e.g., before
FoldablePropagation or
+ * ConstantFolding rewrite predicates).
+ */
+ def preOperatorOptimizationRules: Seq[Rule[LogicalPlan]] = Nil
+
/**
* Override to provide additional rules for early projection and filter
pushdown to scans.
*/
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/SparkSessionExtensions.scala
b/sql/core/src/main/scala/org/apache/spark/sql/SparkSessionExtensions.scala
index 56bff0acc352..5de29732834a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/SparkSessionExtensions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/SparkSessionExtensions.scala
@@ -309,6 +309,23 @@ class SparkSessionExtensions {
optimizerRules += builder
}
+ private[this] val preOperatorOptimizationRuleBuilders =
mutable.Buffer.empty[RuleBuilder]
+
+ private[sql] def buildPreOperatorOptimizationRules(
+ session: SparkSession): Seq[Rule[LogicalPlan]] = {
+ preOperatorOptimizationRuleBuilders.map(_.apply(session)).toSeq
+ }
+
+ /**
+ * Inject an optimizer `Rule` builder into the [[SparkSession]]. The
injected rules will be
+ * executed in a single pass (Once) before the main operator optimization
fixed-point batch.
+ * Use this for rules that need to observe the plan as it enters the
operator-optimization
+ * fixed point, before built-in rules like FoldablePropagation or
ConstantFolding transform it.
+ */
+ def injectPreOperatorOptimizationRule(builder: RuleBuilder): Unit = {
+ preOperatorOptimizationRuleBuilders += builder
+ }
+
private[this] val preCBORules = mutable.Buffer.empty[RuleBuilder]
private[sql] def buildPreCBORules(session: SparkSession):
Seq[Rule[LogicalPlan]] = {
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala
b/sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala
index 4838e6d60f72..52c6821d0001 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala
@@ -332,6 +332,9 @@ abstract class BaseSessionStateBuilder(
override def extendedOperatorOptimizationRules: Seq[Rule[LogicalPlan]] =
super.extendedOperatorOptimizationRules ++
customOperatorOptimizationRules
+
+ override def preOperatorOptimizationRules: Seq[Rule[LogicalPlan]] =
+ super.preOperatorOptimizationRules ++
customPreOperatorOptimizationRules
}
}
@@ -345,6 +348,16 @@ abstract class BaseSessionStateBuilder(
extensions.buildOptimizerRules(session)
}
+ /**
+ * Custom rules that run in a single pass before the main operator
optimization batch.
+ * Prefer overriding this instead of creating your own Optimizer.
+ *
+ * Note that this may NOT depend on the `optimizer` function.
+ */
+ protected def customPreOperatorOptimizationRules: Seq[Rule[LogicalPlan]] = {
+ extensions.buildPreOperatorOptimizationRules(session)
+ }
+
/**
* Custom early scan push down rules to add to the Optimizer. Prefer
overriding this instead
* of creating your own Optimizer.
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
index 89dd1a629af1..d9b29b757b76 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/SparkSessionExtensionSuite.scala
@@ -132,6 +132,12 @@ class SparkSessionExtensionSuite extends PlanTest with
AdaptiveSparkPlanHelper {
}
}
+ test("SPARK-57194: inject pre operator optimization rule") {
+ withSession(Seq(_.injectPreOperatorOptimizationRule(MyRule))) { session =>
+
assert(session.sessionState.optimizer.preOperatorOptimizationRules.contains(MyRule(session)))
+ }
+ }
+
test("inject spark planner strategy") {
withSession(Seq(_.injectPlannerStrategy(MySparkStrategy))) { session =>
assert(session.sessionState.planner.strategies.contains(MySparkStrategy(session)))
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
