[SPARK-XXXXX][SQL] Reuse AttributeSeq when normalizing a sequence in QueryPlan canonicalization

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala

@@ -862,6 +862,19 @@ object QueryPlan extends PredicateHelper {
     }.canonicalized.asInstanceOf[T]
   }
 
+  /**
+   * Normalizes each expression in `exprs` against `input`. Unlike
+   * `exprs.map(normalizeExpressions(_, input))`, this binds the `AttributeSeq` a single time so its
+   * `exprId`-to-ordinal lookup map is built once for the whole sequence. Passing a bare
+   * `Seq[Attribute]` to the per-expression overload inside a `map` re-applies the implicit
+   * `Seq[Attribute]` => `AttributeSeq` conversion on every element, rebuilding that map each time,
+   * which is `O(exprs.size * input.size)` -- quadratic when canonicalizing a relation's own wide
+   * output. The result is identical to the per-element form.
+   */
+  def normalizeExpressions[T <: Expression](exprs: Seq[T], input: AttributeSeq): Seq[T] = {
+    exprs.map(normalizeExpressions(_, input))
+  }
+
   /**
    * Composes the given predicates into a conjunctive predicate, which is normalized and reordered.
    * Then returns a new sequence of predicates by splitting the conjunctive predicate.

sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala

@@ -22,7 +22,7 @@ import java.util.{Collections, Optional, OptionalLong}
 import org.apache.spark.SparkException
 import org.apache.spark.sql.catalyst.analysis.{MultiInstanceRelation, NamedRelation, TimeTravelSpec}
 import org.apache.spark.sql.catalyst.catalog.{CatalogColumnStat, CatalogStatistics}
-import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, Expression, SortOrder, V2ExpressionUtils}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSeq, AttributeSet, Expression, SortOrder, V2ExpressionUtils}
 import org.apache.spark.sql.catalyst.plans.QueryPlan
 import org.apache.spark.sql.catalyst.plans.logical.{ColumnStat, ExposesMetadataColumns, Histogram, HistogramBin, LeafNode, LogicalPlan, Statistics}
 import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils
@@ -285,20 +285,21 @@ case class DataSourceV2ScanRelation(
   }
 
   override def doCanonicalize(): DataSourceV2ScanRelation = {
+    val outputAttrs: AttributeSeq = output
     this.copy(
       relation = this.relation.copy(
-        output = this.relation.output.map(QueryPlan.normalizeExpressions(_, this.relation.output))
+        output = QueryPlan.normalizeExpressions(this.relation.output, this.relation.output)
       ),
-      output = this.output.map(QueryPlan.normalizeExpressions(_, this.output)),
+      output = QueryPlan.normalizeExpressions(this.output, outputAttrs),
       keyGroupedPartitioning = keyGroupedPartitioning.map(
-        _.map(QueryPlan.normalizeExpressions(_, output))
-      ),
-      ordering = ordering.map(
-        _.map(o => o.copy(child = QueryPlan.normalizeExpressions(o.child, output)))
+        QueryPlan.normalizeExpressions(_, outputAttrs)
       ),
+      ordering = ordering.map { orderings =>
+        orderings.map(o => o.copy(child = QueryPlan.normalizeExpressions(o.child, outputAttrs)))
+      },
       // pushedFilters may reference columns pruned out of `output` (see the field doc), so they are
       // normalized against the relation's full output rather than `output`.
-      pushedFilters = pushedFilters.map(QueryPlan.normalizeExpressions(_, relation.output))
+      pushedFilters = QueryPlan.normalizeExpressions(pushedFilters, relation.output)
     )
   }
 }

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/QueryPlanSuite.scala

@@ -200,4 +200,14 @@ class QueryPlanSuite extends SparkFunSuite {
     assert(normalizedExpressions.map(_.dataType) == Seq(IntegerType, StringType),
       "normalizeExpressions preserves original expression order")
   }
+
+  test("normalizeExpressions(Seq) matches the per-element form") {
+    val id = AttributeReference("id", IntegerType)(ExprId(0))
+    val data = AttributeReference("data", StringType)(ExprId(1))
+    val output: AttributeSeq = Seq(id, data)
+    val exprs: Seq[Expression] = Seq(data, id)
+
+    assert(QueryPlan.normalizeExpressions(exprs, output) ==
+      exprs.map(QueryPlan.normalizeExpressions(_, output)))
+  }
 }

sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala

@@ -288,7 +288,7 @@ case class RowDataSourceScanExec(
   // Don't care about `rdd` and `tableIdentifier`, and `stream` when canonicalizing.
   override def doCanonicalize(): SparkPlan =
     copy(
-      output.map(QueryPlan.normalizeExpressions(_, output)),
+      QueryPlan.normalizeExpressions(output, output),
       rdd = null,
       tableIdentifier = None,
       stream = None)
@@ -962,7 +962,7 @@ case class FileSourceScanExec(
       // remove stream on canonicalization; this is needed for reused shuffle to be effective in
       // self-join
       None,
-      output.map(QueryPlan.normalizeExpressions(_, output)),
+      QueryPlan.normalizeExpressions(output, output),
       requiredSchema,
       QueryPlan.normalizePredicates(
         filterUnusedDynamicPruningExpressions(partitionFilters), output),

sql/core/src/main/scala/org/apache/spark/sql/execution/SubqueryAdaptiveBroadcastExec.scala

@@ -49,7 +49,7 @@ case class SubqueryAdaptiveBroadcastExec(
   }
 
   protected override def doCanonicalize(): SparkPlan = {
-    val keys = buildKeys.map(k => QueryPlan.normalizeExpressions(k, child.output))
+    val keys = QueryPlan.normalizeExpressions(buildKeys, child.output)
     copy(name = "dpp", buildKeys = keys, child = child.canonicalized)(None)
   }
 

sql/core/src/main/scala/org/apache/spark/sql/execution/SubqueryBroadcastExec.scala

@@ -70,7 +70,7 @@ case class SubqueryBroadcastExec(
     "collectTime" -> SQLMetrics.createMetric(sparkContext, "time to collect (ms)"))
 
   override def doCanonicalize(): SparkPlan = {
-    val keys = buildKeys.map(k => QueryPlan.normalizeExpressions(k, child.output))
+    val keys = QueryPlan.normalizeExpressions(buildKeys, child.output)
     SubqueryBroadcastExec("dpp", indices, keys, child.canonicalized)
   }
 

sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala

@@ -669,7 +669,7 @@ case class InMemoryRelation(
   override def innerChildren: Seq[SparkPlan] = Seq(cachedPlan)
 
   override def doCanonicalize(): logical.LogicalPlan =
-    withOutput(output.map(QueryPlan.normalizeExpressions(_, output)))
+    withOutput(QueryPlan.normalizeExpressions(output, output))
 
   @transient val partitionStatistics = new PartitionStatistics(output)
 

sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryTableScanExec.scala

@@ -79,10 +79,12 @@ case class InMemoryTableScanExec(
 
   override def innerChildren: Seq[QueryPlan[_]] = Seq(relation) ++ super.innerChildren
 
-  override def doCanonicalize(): SparkPlan =
-    copy(attributes = attributes.map(QueryPlan.normalizeExpressions(_, relation.output)),
-      predicates = predicates.map(QueryPlan.normalizeExpressions(_, relation.output)),
+  override def doCanonicalize(): SparkPlan = {
+    val relationOutput: AttributeSeq = relation.output
+    copy(attributes = QueryPlan.normalizeExpressions(attributes, relationOutput),
+      predicates = QueryPlan.normalizeExpressions(predicates, relationOutput),
       relation = relation.canonicalized.asInstanceOf[InMemoryRelation])
+  }
 
   override def vectorTypes: Option[Seq[String]] =
     relation.cacheBuilder.serializer.vectorTypes(attributes, conf)

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/LogicalRelation.scala

@@ -50,7 +50,7 @@ case class LogicalRelation(
 
   // Only care about relation when canonicalizing.
   override def doCanonicalize(): LogicalPlan = copy(
-    output = output.map(QueryPlan.normalizeExpressions(_, output)),
+    output = QueryPlan.normalizeExpressions(output, output),
     catalogTable = None)
 
   override def computeStats(): Statistics = {

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala

@@ -101,13 +101,14 @@ case class BatchScanExec(
   }
 
   override def doCanonicalize(): BatchScanExec = {
+    val outputAttrs: AttributeSeq = output
     this.copy(
-      output = output.map(QueryPlan.normalizeExpressions(_, output)),
+      output = QueryPlan.normalizeExpressions(output, outputAttrs),
       runtimeFilters = QueryPlan.normalizePredicates(
         runtimeFilters.filterNot(_ == DynamicPruningExpression(Literal.TrueLiteral)),
-        output),
-      keyGroupedPartitioning = keyGroupedPartitioning.map(_.map(
-        QueryPlan.normalizeExpressions(_, output))))
+        outputAttrs),
+      keyGroupedPartitioning = keyGroupedPartitioning.map(
+        QueryPlan.normalizeExpressions(_, outputAttrs)))
   }
 
   override def simpleString(maxFields: Int): String = {