(spark) branch master updated: [SPARK-59597][SQL] Avoid quadratic CTE attribute restoration

[ur...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

uros-b 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 186c4333e7f0 [SPARK-59597][SQL] Avoid quadratic CTE attribute restoration
186c4333e7f0 is described below

commit 186c4333e7f043937af31ab06a0134299d938598
Author: bhollis-dbx <be...@databricks.com>
AuthorDate: Tue Sep 29 19:51:14 2026 +0000

    [SPARK-59597][SQL] Avoid quadratic CTE attribute restoration
    
    ### What changes were proposed in this pull request?
    
    Use `AttributeMap` when restoring attributes from CTE references to CTE definitions in `PushdownPredicatesAndPruneColumnsForCTEDef`. This replaces a linear semantic search per attribute with expression-ID lookup.
    
    ### Why are the changes needed?
    
    The current implementation linearly searches the CTE attribute mapping for every restored attribute, making planning quadratic in CTE width. The cost reaches tens of milliseconds per rule application for wide CTEs and is repeated for CTE references and optimizer iterations.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    Ran `PushdownPredicatesForCTEDefStalenessSuite`. All 7 tests passed.
    
    Local benchmark output for the full rule:
    
    | Columns | Before | After | Change |
    |---:|---:|---:|---:|
    | 1 | 2.38 us | 2.48 us | 4% slower |
    | 2 | 2.85 us | 2.92 us | 2% slower |
    | 5 | 5.02 us | 4.72 us | 6% faster |
    | 10 | 8.15 us | 7.21 us | 12% faster |
    | 20 | 14.78 us | 12.41 us | 16% faster |
    | 50 | 41.36 us | 28.96 us | 30% faster |
    | 100 | 100.11 us | 54.71 us | 45% faster |
    | 500 | 1.65 ms | 0.27 ms | 6.1x faster |
    | 1,000 | 6.53 ms | 0.58 ms | 11.2x faster |
    | 2,000 | 26.54 ms | 1.33 ms | 19.9x faster |
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: OpenAI Codex (GPT-5)
    
    Closes #58873 from bhollis-dbx/optimize-cte-restore-attrs.
    
    Authored-by: bhollis-dbx <be...@databricks.com>
    Signed-off-by: Uros Bojanic <22...@users.noreply.github.com>
---
 .../optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala    | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala
index b7e875554400..43b6845aa01f 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/PushdownPredicatesAndPruneColumnsForCTEDef.scala
@@ -19,13 +19,12 @@ package org.apache.spark.sql.catalyst.optimizer
 
 import scala.collection.mutable
 
-import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeSet}
+import org.apache.spark.sql.catalyst.expressions.{And, Attribute, AttributeMap, AttributeSet}
 import org.apache.spark.sql.catalyst.expressions.{Expression, Literal, Or, PredicateHelper, SubqueryExpression}
 import org.apache.spark.sql.catalyst.planning.PhysicalOperation
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.trees.TreePattern.CTE
-import org.apache.spark.util.collection.Utils
 
 /**
  * Infer predicates and column pruning for [[CTERelationDef]] from its reference points, and push
@@ -50,8 +49,7 @@ object PushdownPredicatesAndPruneColumnsForCTEDef extends Rule[LogicalPlan] with
       input: Seq[Expression],
       mapping: Map[Attribute, Expression]): Seq[Expression] = {
     input.map(e => e.transform {
-      case a: Attribute =>
-        mapping.keys.find(_.semanticEquals(a)).map(mapping).getOrElse(a)
+      case a: Attribute => mapping.getOrElse(a, a)
     })
   }
 
@@ -75,7 +73,7 @@ object PushdownPredicatesAndPruneColumnsForCTEDef extends Rule[LogicalPlan] with
 
       case PhysicalOperation(projects, predicates, ref: CTERelationRef) =>
         val (cteDef, precedence, preds, attrs) = cteMap(ref.cteId)
-        val attrMapping = Utils.toMap(ref.output, cteDef.output)
+        val attrMapping = AttributeMap(ref.output.zip(cteDef.output))
         val newPredicates = if (isTruePredicate(preds)) {
           preds
         } else {


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org