[SPARK-59261][SQL] Memoize outputPartitioning on the projection and SPJ grouping nodes

[peter-toth]

What changes were proposed in this pull request?

PartitioningPreservingUnaryExecNode.outputPartitioning and GroupPartitionsExec.outputPartitioning become lazy vals. 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 LazyLists. 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)