This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 0b4d79395e3e [SPARK-56877][SQL][FOLLOWUP] Make PartitioningCollection
invariant check O(1) per nesting level
0b4d79395e3e is described below
commit 0b4d79395e3eff4b8fcc4007bb7f79c4527a29ea
Author: Wenchen Fan <[email protected]>
AuthorDate: Thu Jun 11 10:44:31 2026 -0700
[SPARK-56877][SQL][FOLLOWUP] Make PartitioningCollection invariant check
O(1) per nesting level
### What changes were proposed in this pull request?
Make the `KeyedPartitioning` invariant enforcement added in #55901
(SPARK-56877) cheap for deeply nested `PartitioningCollection`s:
- `checkKeyedPartitioningInvariant()` no longer walks the entire
partitioning tree with `TreeNode.foreach` on every construction. Each
collection now exposes a cached representative (`firstKeyedPartitioning`);
since every nested collection already enforced the invariant on its own
construction, comparing one representative per direct member validates the
whole subtree by induction. The check is now O(partitionings.size) instead of
O(subtree size).
- `PartitioningCollection.fromPartitionings` no longer rebuilds nested
collections that are already consistent. Using the same representative, a
nested collection with no `KeyedPartitioning` in its subtree, or whose
canonical `partitionKeys` reference already matches, is returned as-is in O(1).
It only recurses and rebuilds when interning is actually needed.
The invariant itself is unchanged and still enforced in the constructor.
### Why are the changes needed?
SPARK-56877 caused a planning-time regression for plans that chain many
shuffle joins on the same key. For an inner join,
`ShuffledJoin.outputPartitioning` wraps the children's partitionings in a
`PartitioningCollection`, so a chain of N same-key joins nests collections N
levels deep, and `outputPartitioning` is a `def` recomputed on every access.
With the constructor running a full-subtree walk and `fromPartitionings`
rebuilding every nested level (re-triggering the walk at each leve [...]
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
New unit tests in `DistributionSuite`:
- `fromPartitionings` returns already-consistent nested collections
reference-equal (guards against reintroducing the rebuild),
- `fromPartitionings` still interns
structurally-equal-but-reference-distinct `partitionKeys` across nesting,
- the constructor still rejects `partitionKeys` reference mismatches and
expression arity mismatches through nesting.
Existing suites covering the invariant and interning behavior pass:
`DistributionSuite`, `EnsureRequirementsSuite`,
`ProjectedOrderingAndPartitioningSuite`, `GroupPartitionsExecSuite`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Fable 5)
Closes #56411 from cloud-fan/fix-partitioning-collection-invariant-perf.
Authored-by: Wenchen Fan <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
(cherry picked from commit 302ba67264fe56008105c096ec1f06eb6645b81c)
Signed-off-by: Wenchen Fan <[email protected]>
---
.../sql/catalyst/plans/physical/partitioning.scala | 60 ++++++++++++++++++----
.../spark/sql/catalyst/DistributionSuite.scala | 56 ++++++++++++++++++++
2 files changed, 106 insertions(+), 10 deletions(-)
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
index aeacdaec7a8d..d2bb12d2053a 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
@@ -796,21 +796,44 @@ case class PartitioningCollection(partitionings:
Seq[Partitioning])
checkKeyedPartitioningInvariant()
+ /**
+ * First [[KeyedPartitioning]] reachable from this collection through direct
members or nested
+ * collections, if any. Since every collection validates the invariant on
construction, this
+ * single representative stands for all [[KeyedPartitioning]]s in the
subtree: they all share
+ * its `partitionKeys` reference and expression arity. The invariant check
forces this lazy val
+ * during construction, so it is only recomputed after deserialization.
+ */
+ @transient private[physical] lazy val firstKeyedPartitioning:
Option[KeyedPartitioning] =
+ partitionings.view.map {
+ case k: KeyedPartitioning => Some(k)
+ case pc: PartitioningCollection => pc.firstKeyedPartitioning
+ case _ => None
+ }.collectFirst { case Some(k) => k }
+
+ /**
+ * Nested collections already enforced the invariant on their own
construction, so comparing one
+ * representative per direct member against [[firstKeyedPartitioning]]
validates the whole
+ * subtree without walking it. Keeping this check O(partitionings.size)
matters: join
+ * `outputPartitioning` builds these collections afresh on every call, and
plans chaining many
+ * same-key joins nest them linearly deep.
+ */
private def checkKeyedPartitioningInvariant(): Unit = {
- var first: KeyedPartitioning = null
- foreach {
- case k: KeyedPartitioning =>
- if (first == null) {
- first = k
- } else {
- require(k.expressions.length == first.expressions.length,
+ firstKeyedPartitioning.foreach { first =>
+ partitionings.foreach { p =>
+ val representative = p match {
+ case k: KeyedPartitioning => k
+ case pc: PartitioningCollection => pc.firstKeyedPartitioning.orNull
+ case _ => null
+ }
+ if (representative != null && (representative ne first)) {
+ require(representative.expressions.length ==
first.expressions.length,
"All KeyedPartitionings in a PartitioningCollection must have
matching expression " +
"arity")
- require(k.partitionKeys eq first.partitionKeys,
+ require(representative.partitionKeys eq first.partitionKeys,
"All KeyedPartitionings in a PartitioningCollection must share the
same " +
"partitionKeys reference")
}
- case _ =>
+ }
}
}
@@ -868,7 +891,24 @@ object PartitioningCollection {
} else {
k
}
- case pc: PartitioningCollection => new
PartitioningCollection(pc.partitionings.map(intern))
+ case pc: PartitioningCollection =>
+ pc.firstKeyedPartitioning match {
+ // No KeyedPartitioning anywhere in this subtree: nothing to intern.
Returning the
+ // collection as-is keeps repeated outputPartitioning computations
over deeply nested
+ // collections (e.g. chains of same-key joins) O(1) per level.
+ case None => pc
+ case Some(representative) if canonicalKeys == null =>
+ canonicalKeys = representative.partitionKeys
+ pc
+ // The collection's own invariant guarantees all its
KeyedPartitionings share the
+ // representative's `partitionKeys` reference, so reference-equality
of the
+ // representative's keys means the whole subtree is already interned.
+ case Some(representative) if representative.partitionKeys eq
canonicalKeys => pc
+ case Some(representative) =>
+ require(representative.partitionKeys == canonicalKeys,
+ "All KeyedPartitionings in a PartitioningCollection must have
equal partitionKeys")
+ new PartitioningCollection(pc.partitionings.map(intern))
+ }
case other => other
}
new PartitioningCollection(partitionings.map(intern))
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/DistributionSuite.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/DistributionSuite.scala
index 88faf7010318..2e9e1270b3fc 100644
---
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/DistributionSuite.scala
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/DistributionSuite.scala
@@ -410,4 +410,60 @@ class DistributionSuite extends SparkFunSuite {
assert(groupedKP.isGrouped)
checkSatisfied(groupedKP, ClusteredDistribution(Seq(x)), true)
}
+
+ test("SPARK-56877: fromPartitionings reuses already-consistent nested
collections") {
+ val x = AttributeReference("x", IntegerType)()
+ val y = AttributeReference("y", IntegerType)()
+
+ // No KeyedPartitioning anywhere in the subtree: the nested collection is
returned as-is.
+ // Rebuilding it would make outputPartitioning of deeply nested
collections (e.g. chains of
+ // same-key shuffle joins) quadratic in the nesting depth.
+ val hashCollection = PartitioningCollection.fromPartitionings(
+ Seq(HashPartitioning(Seq(x), 10), HashPartitioning(Seq(y), 10)))
+ val wrapped = PartitioningCollection.fromPartitionings(
+ Seq(hashCollection, HashPartitioning(Seq(y), 10)))
+ assert(wrapped.partitionings.head eq hashCollection)
+
+ // KeyedPartitionings already share the canonical partitionKeys reference:
also as-is.
+ val kpX = KeyedPartitioning(Seq(x), Seq(InternalRow(1), InternalRow(2),
InternalRow(3)))
+ val kpY = kpX.copy(expressions = Seq(y))
+ val keyedCollection = PartitioningCollection.fromPartitionings(Seq(kpX,
kpY))
+ val keyedWrapped =
PartitioningCollection.fromPartitionings(Seq(keyedCollection, kpX))
+ assert(keyedWrapped.partitionings.head eq keyedCollection)
+ }
+
+ test("SPARK-56877: fromPartitionings interns partitionKeys across nested
collections") {
+ val x = AttributeReference("x", IntegerType)()
+ val y = AttributeReference("y", IntegerType)()
+
+ val kpX = KeyedPartitioning(Seq(x), Seq(InternalRow(1), InternalRow(2)))
+ val nested = PartitioningCollection.fromPartitionings(Seq(kpX))
+ // Structurally equal but reference-distinct partitionKeys.
+ val kpY = KeyedPartitioning(Seq(y), Seq(InternalRow(1), InternalRow(2)))
+ assert(kpX.partitionKeys ne kpY.partitionKeys)
+
+ val combined = PartitioningCollection.fromPartitionings(Seq(nested, kpY))
+ val interned = combined.partitionings.last.asInstanceOf[KeyedPartitioning]
+ assert(interned.partitionKeys eq kpX.partitionKeys)
+ }
+
+ test("SPARK-56877: PartitioningCollection enforces the invariant through
nesting") {
+ val x = AttributeReference("x", IntegerType)()
+ val y = AttributeReference("y", IntegerType)()
+
+ val kpX = KeyedPartitioning(Seq(x), Seq(InternalRow(1), InternalRow(2)))
+ val nested = PartitioningCollection.fromPartitionings(Seq(kpX))
+
+ val kpY = KeyedPartitioning(Seq(y), Seq(InternalRow(1), InternalRow(2)))
+ val refMismatch = intercept[IllegalArgumentException] {
+ PartitioningCollection(Seq(nested, kpY))
+ }
+ assert(refMismatch.getMessage.contains("share the same partitionKeys
reference"))
+
+ val kpXY = KeyedPartitioning(Seq(x, y), Seq(InternalRow(1, 1),
InternalRow(2, 2)))
+ val arityMismatch = intercept[IllegalArgumentException] {
+ PartitioningCollection(Seq(nested, kpXY))
+ }
+ assert(arityMismatch.getMessage.contains("matching expression arity"))
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
