From c53242d6bd88b93d4b1c1d919dac93057f06007f Mon Sep 17 00:00:00 2001 From: Ben Hollis Date: Wed, 16 Sep 2026 22:13:44 +0000 Subject: [PATCH] [SQL] Avoid quadratic CTE attribute restoration --- .../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 d33384152e8b9..b2e256d895862 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 {