From e01f54455fa836e6dfbb722f22e61df1de3b0811 Mon Sep 17 00:00:00 2001 From: Peter Toth Date: Fri, 4 Sep 2026 20:06:20 +0200 Subject: [PATCH] [SPARK-59261][SQL] Memoize outputPartitioning on the projection and SPJ grouping nodes ### What changes were proposed in this pull request? `PartitioningPreservingUnaryExecNode.outputPartitioning` and `GroupPartitionsExec.outputPartitioning` become `lazy val`s. The projection body also binds `child.outputPartitioning` to a local, because it read it twice. ### Why are the changes needed? Both recompute a non-trivial value on every call, and the planner asks a node for its partitioning many times. Counted by instrumenting the two bodies and running `KeyGroupedPartitioningSuite`: PartitioningPreservingUnaryExecNode 23,512 -> 3,663 GroupPartitionsExec 11,110 -> 1,326 Both figures are this PR against its own base, with the two changes applied together. `PartitioningPreservingUnaryExecNode` is the expensive one. For every partitioning it projects each expression through the output aliases and deduplicates the results by `canonicalized`. For a `KeyedPartitioning` child it also allocates an `ExpressionSet` per key position and cross-products the per-position alternatives through `LazyList`s. It is the biggest single caller of the V2 scan's `outputPartitioning` too. On the same instrumented run, 19,105 of that method's 33,051 reads arrive through it, mostly as `Project -> Filter -> scan`. `GroupPartitionsExec` rebuilds every `KeyedPartitioning` in the child's partitioning through `p.transform`, on top of the `grouping` it has already memoized. ### Why memoizing is safe Neither body reads a live config, so neither freezes one. `PartitioningPreservingUnaryExecNode` reads `child.outputPartitioning`, `outputExpressions` and `aliasCandidateLimit`, and that last one is already a `protected val` evaluated at node construction. `GroupPartitionsExec.outputPartitioning` reads `child.outputPartitioning` and its own `grouping`, and no config at all. The planner already treats `outputPartitioning` as a property of the node rather than a question to re-ask. `ValidateRequirements` reads one child's partitioning twice and compares specs built from the two reads. `EnsureRequirements` reads it repeatedly and threads the results between the reads, and at `EnsureRequirements.scala:335` it is the pruning predicate of a `multiTransformDownWithPruning`, re-evaluated per generated alternative. A child's partitioning is not fixed in every case. `InMemoryTableScanExec.outputPartitioning` reports `UnknownPartitioning(0)` while its AQE cached plan is not final, and sharpens once it is. Memoizing here is still right, for two reasons. In the executed plan the ancestor instances are fresh, because `CollapseCodegenStages` and `ApplyColumnarRulesAndInsertTransitions` insert nodes and every ancestor above an insertion is copied, so the memo is taken no earlier than the old `def` was first called. Where an instance does survive such a change, the pinned answer is the one `EnsureRequirements` planned against, which is the answer a later reader should see. The memo retains one `Partitioning` per node for the plan's lifetime, holding up to `aliasCandidateLimit` alternatives. That is a deliberate trade against the recomputations counted above. An identity key projection retains nothing new, because `KeyedPartitioning.project` returns `this` when it drops no position. ### Why a plain `lazy val` `BroadcastHashJoinExec`, `AQEShuffleReadExec` and `FileSourceScanExec` already override `outputPartitioning` as a plain `lazy val`, and all three postdate SPARK-50705, which added `BestEffortLazyVal` for `QueryPlan` members that a tree walk can reach from two directions at once. That hazard needs two lock orders. Both bodies here only descend into `child`, and a physical plan node holds no parent pointer that any `outputPartitioning` implementation follows, so every lock order is parent then child. ### Why not the `outputOrdering` twins `AliasAwareQueryOutputOrdering.outputOrdering` runs the same alias machinery per call, but it is declared `final` in catalyst and is shared with logical plans, where the optimizer creates and discards nodes constantly. That is a wider change than this one. `GroupPartitionsExec.outputOrdering` reads `conf.v2BucketingPreserveKeyOrderingOnCoalesceEnabled`, so memoizing it would change when that config is read. That is a separate question. ### Does this PR introduce _any_ user-facing change? No. `PartitioningPreservingUnaryExecNode.outputPartitioning` was `final override def` and becomes `final override lazy val`. Scala forbids a `def` overriding a `lazy val`, but this member is `final`, so no subclass could have overridden it either way. `GroupPartitionsExec.outputPartitioning` was not final. An out-of-tree subclass overriding it with a `def` would need to become a `lazy val`. Binary compatibility is unaffected. ### How was this patch tested? Existing suites, since the values are unchanged. `KeyGroupedPartitioningSuite`, `GroupPartitionsExecSuite`, `ProjectedOrderingAndPartitioningSuite`, `EnsureRequirementsSuite` and `PlannerSuite`, 311 tests. One of them covers the config claim directly. `ProjectedOrderingAndPartitioningSuite`'s "SPARK-46367: narrowing projection with duplicate keys requires allowKeysSubsetOfPartitionKeys to satisfy ClusteredDistribution" reads the same `ProjectExec` instance with `V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS` off and then on, and asserts a different answer each time. It passes with the partitioning memoized, because that config is read by `mayGroupToSatisfy` on the returned value rather than by the body. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 5) --- .../sql/execution/AliasAwareOutputExpression.scala | 11 ++++++++--- .../datasources/v2/GroupPartitionsExec.scala | 5 ++++- 2 files changed, 12 insertions(+), 4 deletions(-) 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