diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala index ed4e889aa5d98..fea1df4aec8e2 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/QueryPlan.scala +++ b/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. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala index da2b3a54d74c3..362081a8be20e 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala +++ b/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) ) } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/QueryPlanSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/QueryPlanSuite.scala index 68c0e354950f3..359506658374a 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/plans/QueryPlanSuite.scala +++ b/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))) + } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala index a727ccf565063..28db38b2260c3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/DataSourceScanExec.scala +++ b/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), diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SubqueryAdaptiveBroadcastExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SubqueryAdaptiveBroadcastExec.scala index c435f06ffddc4..8665859f79cfb 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SubqueryAdaptiveBroadcastExec.scala +++ b/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) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SubqueryBroadcastExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SubqueryBroadcastExec.scala index 9e7c1193c8ae0..7d8ae074422c4 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SubqueryBroadcastExec.scala +++ b/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) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala index e6bf65ec89e6b..7ab507f5b8265 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala +++ b/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) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryTableScanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryTableScanExec.scala index cbd60804b27e8..cf718070c399d 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryTableScanExec.scala +++ b/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) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/LogicalRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/LogicalRelation.scala index 3c65ef139ea0a..56d88d44526ef 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/LogicalRelation.scala +++ b/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 = { diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala index 2e3394ac2a082..9aa982ce72038 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala +++ b/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 = {