diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala index 6fca19fbf0a90..7d0afa8bb676a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala @@ -29,9 +29,14 @@ import org.apache.spark.sql.catalyst.trees.MultiTransform */ trait PartitioningPreservingUnaryExecNode extends UnaryExecNode with AliasAwareOutputExpression { - final override def outputPartitioning: Partitioning = { + // A `lazy val` because the planner asks a node for its partitioning many times, and this body + // projects every partitioning expression through the output aliases, and for a + // `KeyedPartitioning` child also builds an `ExpressionSet` per key position and cross-products + // the per-position alternatives. Read no live config here, or memoizing would freeze it. + @transient final override lazy val outputPartitioning: Partitioning = { + val childPartitioning = child.outputPartitioning val (keyedPartitionings, otherPartitionings) = - PartitioningCollection.flatten(child.outputPartitioning) + PartitioningCollection.flatten(childPartitioning) .partition(_.isInstanceOf[KeyedPartitioning]) val projectedKPs = @@ -44,7 +49,7 @@ trait PartitioningPreservingUnaryExecNode extends UnaryExecNode // deep projection chain that nesting overflows the stack when the partitioning is later // serialized or deeply traversed. (projectedKPs ++ projectedOthers).take(aliasCandidateLimit).toList match { - case Seq() => UnknownPartitioning(child.outputPartitioning.numPartitions) + case Seq() => UnknownPartitioning(childPartitioning.numPartitions) case Seq(p) => p case ps => PartitioningCollection.fromPartitionings(ps) } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala index eb0142c5a01bd..c785218f49989 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala @@ -67,7 +67,10 @@ case class GroupPartitionsExec( @transient enableSortedMerge: Boolean = false ) extends UnaryExecNode { - override def outputPartitioning: Partitioning = { + // A `lazy val` because the planner asks a node for its partitioning many times, and this body + // rebuilds every `KeyedPartitioning` in the child's on top of the already memoized `grouping`. + // Read no live config here, or memoizing would freeze it. + @transient override lazy val outputPartitioning: Partitioning = { child.outputPartitioning match { case p: Partitioning with Expression => // There can be multiple `KeyedPartitioning`s in an output partitioning of a join, but they