From 6136919314809cd61473fb045d0521c202b41fff Mon Sep 17 00:00:00 2001 From: Peter Toth Date: Thu, 3 Sep 2026 14:45:52 +0200 Subject: [PATCH 1/3] [SPARK-59080][SQL] Pick one ShuffleSpecCollection member for the SPJ pushdown and the re-shuffle Under `spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys`, `KeyedPartitioning.createShuffleSpec` projects each member of a `PartitioningCollection` onto its own subset of the operation keys. The members of the resulting `ShuffleSpecCollection` can therefore end up with different partition counts. `EnsureRequirements` asked the collection itself for a shuffle template, and `ShuffleSpecCollection.createPartitioning` throws `expected all specs in the collection to have the same number of partitions`. The collection cannot answer that question. `isCompatibleWith` succeeds when *any* member matches, so the collection alone never said which member the two sides agreed on, and reading `specs.head` took whichever the alias cross-product enumerated first. `EnsureRequirements` now resolves that member once, preferring the finest when several qualify, and uses it to build the re-shuffled child's partitioning. The `require` stays as a guard on a method that no longer has a production caller. The `joinKeyPositions` pushed into a compatible child now come from that child's own matching member. They index into the child's own partition expressions, so the best spec's positions were only right when the best spec was that child's. No query reaches the wrong case today: a join is handled by `checkKeyGroupCompatible`, which already pushes each side's own positions, and a cogroup's grouping key is synthesized so neither side stays keyed. --- .../spark/sql/catalyst/ShuffleSpecSuite.scala | 32 ++++++++- .../exchange/EnsureRequirements.scala | 47 +++++++++---- .../KeyGroupedPartitioningSuite.scala | 51 ++++++++++++++ .../exchange/EnsureRequirementsSuite.scala | 70 +++++++++++++++++++ 4 files changed, 184 insertions(+), 16 deletions(-) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala index b6cf5dec1f5f7..496b408520fbc 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst import org.apache.spark.{SparkFunSuite, SparkUnsupportedOperationException} import org.apache.spark.sql.catalyst.dsl.expressions._ -import org.apache.spark.sql.catalyst.expressions.{Attribute, DirectShufflePartitionID, Expression, TransformExpression} +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, DirectShufflePartitionID, Expression, TransformExpression} import org.apache.spark.sql.catalyst.plans.SQLHelper import org.apache.spark.sql.catalyst.plans.physical._ import org.apache.spark.sql.connector.catalog.functions.ScalarFunction @@ -690,4 +690,34 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper { expected = false ) } + + test("SPARK-59080: a collection whose members cover different key subsets disagrees") { + val id = AttributeReference("id", IntegerType)() + val t1 = AttributeReference("t1", IntegerType)() + val t2 = AttributeReference("t2", IntegerType)() + val keys = Seq(InternalRow(1, 1), InternalRow(1, 2), InternalRow(2, 1)) + + // The shape an alias cross-product produces: same arity, same keys, different expressions. The + // operation clusters on (id, t1), so the first member projects onto both positions and keeps + // three partitions, while the second matches only `id` and keeps two. + val collection = PartitioningCollection.fromPartitionings(Seq( + KeyedPartitioning(Seq(id, t1), keys), + KeyedPartitioning(Seq(id, t2), keys))) + + withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { + val spec = collection.createShuffleSpec(ClusteredDistribution(Seq(id, t1))) + .asInstanceOf[ShuffleSpecCollection] + + // The disagreement is kept rather than resolved here. Every member has to stay for + // `isCompatibleWith`, which answers for any of them, and the collection cannot know which one + // the other side matched. `EnsureRequirements` resolves that and asks the member, not the + // collection. + assert(spec.specs.map(_.numPartitions).toSet === Set(3, 2)) + assert(spec.isCompatibleWith(spec), "every member stays available for matching") + + // So asking the collection for a single answer is the caller's mistake, and it says so. + val e = intercept[IllegalArgumentException](spec.createPartitioning(Seq(id, t1))) + assert(e.getMessage.contains("expected all specs in the collection to have the same number")) + } + } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala index 2a33280861da0..65f79ddac099a 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala @@ -17,7 +17,6 @@ package org.apache.spark.sql.execution.exchange -import scala.annotation.tailrec import scala.collection.immutable.BitSet import scala.collection.mutable import scala.collection.mutable.ArrayBuffer @@ -247,30 +246,50 @@ case class EnsureRequirements( } } + // A `ShuffleSpecCollection` answers `isCompatibleWith` if *any* of its members does, so the + // collection alone does not say which member the sides agreed on. The projection pushed into + // a compatible child and the partitioning built for a re-shuffled child both have to come + // from one member, otherwise the sides end up grouped on different keys, or on a key set the + // child does not even have. Pick that member once, preferring the finest when several + // qualify. Only the branch that shuffles a child reads these, hence `lazy`. + lazy val matchedIndexes = bestSpecOpt.toSeq.flatMap { best => + childrenIndexes.filter(i => best.isCompatibleWith(specs(i))) + } + lazy val bestMemberOpt = bestSpecOpt.flatMap { best => + val matchedMembers = matchedIndexes.map(i => flattenSpec(specs(i))) + // No member serving every matched child means there is no layout to align them on, so they + // all take the ordinary shuffle. That needs three or more clustered children, since with + // two the member that reported the match serves both, and no operator has three today. + flattenSpec(best) + .filter(m => matchedMembers.forall(_.exists(m.isCompatibleWith))) + .maxByOption(_.numPartitions) + } + children = children.zip(requiredChildDistributions).zipWithIndex.map { case ((child, _), idx) if areChildrenCompatible || !childrenIndexes.contains(idx) => child case ((child, dist), idx) => - if (bestSpecOpt.isDefined && bestSpecOpt.get.isCompatibleWith(specs(idx))) { - // If the child's partitioning is a `PartitioningCollection`, its spec is a - // `ShuffleSpecCollection` whose `createPartitioning` delegates to the head spec, - // so unwrap to the head spec to stay aligned with the re-shuffled side below. - unwrapSpecCollection(bestSpecOpt.get) match { + if (bestMemberOpt.isDefined && matchedIndexes.contains(idx)) { + // The positions come from this child's own matching member, since they index into its + // own partition expressions -- the chosen best member only says which member of it the + // two sides agreed on. + val bestMember = bestMemberOpt.get + flattenSpec(specs(idx)).find(bestMember.isCompatibleWith) match { // If `areChildrenCompatible` is false, we can still perform SPJ // by shuffling the other side based on join keys (see the else case below). // Hence we need to ensure that after this call, the outputPartitioning of the // partitioned side's BatchScanExec is grouped by join keys to match, // and we do that by pushing down the join keys - case KeyedShuffleSpec(_, _, Some(joinKeyPositions)) => + case Some(KeyedShuffleSpec(_, _, Some(joinKeyPositions))) => withJoinKeyPositions(child, joinKeyPositions) case _ => child } } else { - val newPartitioning = bestSpecOpt.map { bestSpec => + val newPartitioning = bestMemberOpt.map { bestMember => // Use the best spec to create a new partitioning to re-shuffle this child val clustering = dist.asInstanceOf[ClusteredDistribution].clustering - bestSpec.createPartitioning(clustering) + bestMember.createPartitioning(clustering) }.getOrElse { // No best spec available, so we create default partitioning from the required // distribution @@ -779,12 +798,10 @@ case class EnsureRequirements( } } - // Unwraps a `ShuffleSpecCollection` (possibly nested) to the spec that its - // `createPartitioning` delegates to, i.e. the head spec. - @tailrec - private def unwrapSpecCollection(spec: ShuffleSpec): ShuffleSpec = spec match { - case ShuffleSpecCollection(specs) => unwrapSpecCollection(specs.head) - case other => other + // Flattens a (possibly nested) `ShuffleSpecCollection` into its member specs. + private def flattenSpec(spec: ShuffleSpec): Seq[ShuffleSpec] = spec match { + case ShuffleSpecCollection(specs) => specs.flatMap(flattenSpec) + case other => Seq(other) } /** diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala index 6a911ce652615..c700ce6cbd3d3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala @@ -3201,6 +3201,57 @@ class KeyGroupedPartitioningSuite } } + test("SPARK-59080: both sides of the join land on the same collection member") { + // `arrive_time` is selected twice under two aliases, so the projected partitioning is a + // `PartitioningCollection` whose members cover different numbers of the join keys: one covers + // (id, t1), another only id. Each member's spec is projected onto its own subset, so the specs + // disagree on `numPartitions`, and asking the collection for one partitioning fails with + // "expected all specs in the collection to have the same number of partitions". + // + // `EnsureRequirements` now resolves the member the two sides agreed on before it asks, so the + // keyed side is grouped on both join keys and the shuffled side is laid out on those same keys. + val items_partitions = Array(identity("id"), identity("arrive_time")) + createTable(items, itemsColumns, items_partitions) + + sql(s"INSERT INTO testcat.ns.$items VALUES " + + "(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + + "(1, 'ab', 30.0, cast('2020-01-02' as timestamp)), " + + "(3, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + + "(4, 'cc', 15.5, cast('2020-02-01' as timestamp))") + + createTable(purchases, purchasesColumns, Array.empty) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + "(1, 42.0, cast('2020-01-01' as timestamp)), " + + "(1, 89.0, cast('2020-01-02' as timestamp)), " + + "(3, 19.5, cast('2020-01-01' as timestamp)), " + + "(5, 26.0, cast('2023-01-01' as timestamp))") + + withSQLConf( + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { + val df = sql( + s""" + |${selectWithMergeJoinHint("i", "p")} + |id, t1, t2, i.price AS purchase_price, p.price AS sale_price + |FROM (SELECT id, arrive_time AS t1, arrive_time AS t2, price FROM testcat.ns.$items) i + |JOIN testcat.ns.$purchases p ON i.id = p.item_id AND i.t1 = p.time + |""".stripMargin) + val plan = df.queryExecution.executedPlan + val positions = collectAllGroupPartitions(plan).flatMap(_.joinKeyPositions) + assert(positions === Seq(Seq(0, 1)), + "the keyed side must be grouped on both join keys, the finest granularity available") + assert(collectAllShuffles(plan).size == 1, "only the unpartitioned side shuffles") + checkAnswer(df, Seq( + Row(1, java.sql.Timestamp.valueOf("2020-01-01 00:00:00"), + java.sql.Timestamp.valueOf("2020-01-01 00:00:00"), 40.0, 42.0), + Row(1, java.sql.Timestamp.valueOf("2020-01-02 00:00:00"), + java.sql.Timestamp.valueOf("2020-01-02 00:00:00"), 30.0, 89.0), + Row(3, java.sql.Timestamp.valueOf("2020-01-01 00:00:00"), + java.sql.Timestamp.valueOf("2020-01-01 00:00:00"), 10.0, 19.5))) + } + } + test("SPARK-59025: shuffle one side and join keys are less than partition keys " + "when the keyed side reports a PartitioningCollection") { val items_partitions = Array(identity("id"), identity("name")) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala index 1d877f169605c..d29b0a8797c24 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala @@ -1421,6 +1421,38 @@ class EnsureRequirementsSuite extends SharedSparkSession { TransformExpression(DaysFunction, Seq(expr)) } + test("SPARK-59080: the re-shuffled side lands on the member the keyed side was matched on") { + val id = AttributeReference("id", IntegerType)() + val t1 = AttributeReference("t1", IntegerType)() + val t2 = AttributeReference("t2", IntegerType)() + // Same arity and the same keys, different expressions: what an alias cross-product produces + // when one column is selected twice. Clustering on (id, t1), the first member covers `id` only + // and keeps two partitions, the second covers both positions and keeps three. The coarse + // member comes first so that reading the collection's head would pick the wrong one. + val keys = Seq(InternalRow(1, 1), InternalRow(1, 2), InternalRow(2, 1)) + val keyed = new DummySparkPlanWithBatchScanChild( + outputPartitioning = PartitioningCollection.fromPartitionings(Seq( + KeyedPartitioning(Seq(id, t2), keys), + KeyedPartitioning(Seq(id, t1), keys)))) + val unpartitioned = DummySparkPlan(outputPartitioning = UnknownPartitioning(0)) + + withSQLConf( + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { + val smj = SortMergeJoinExec(Seq(id, t1), Seq(id, t1), Inner, None, keyed, unpartitioned) + val planned = EnsureRequirements.apply(smj).asInstanceOf[SortMergeJoinExec] + + // The keyed side is grouped on both join keys, which is the finest member. + assert(groupPartitionsNodes(planned.left).map(_.joinKeyPositions) === Seq(Some(Seq(0, 1))), + "the keyed side must be grouped on the member covering both join keys") + // And the shuffled side lands on that same member. Reading the collection instead would take + // whichever member came first and could put the two sides on different key sets. + val shuffles = planned.right.collect { case s: ShuffleExchangeExec => s } + assert(shuffles.map(_.outputPartitioning.numPartitions) === Seq(3), + "the re-shuffled side must land on the same member, not on whichever came first") + } + } + private class DummySparkPlanWithBatchScanChild(outputPartitioning: Partitioning) extends DummySparkPlan( children = Seq(BatchScanExec(Seq.empty, null, Seq.empty, table = null)), @@ -1429,6 +1461,44 @@ class EnsureRequirementsSuite extends SharedSparkSession { requiredChildOrdering = Seq(Seq.empty) ) + test("SPARK-59080: pushed-down positions index into the child's own partition expressions") { + val id = AttributeReference("id", IntegerType)() + val t1 = AttributeReference("t1", IntegerType)() + val other = AttributeReference("other", IntegerType)() + // Both children carry the same (id, t1) key set, but the second declares a leading partition + // expression the distribution does not cluster on. So the keys it projects onto sit at + // positions 1 and 2, while the first child's sit at 0 and 1. + // + // No query reaches two keyed children here: a join is handled by `checkKeyGroupCompatible`, + // which pushes each side's own positions, and a cogroup's grouping key is synthesized so + // neither side stays keyed. The test pins the invariant rather than reproducing a query. + val keys = Seq(InternalRow(1, 1), InternalRow(1, 2), InternalRow(2, 1)) + val paddedKeys = Seq(InternalRow(9, 1, 1), InternalRow(9, 1, 2), InternalRow(9, 2, 1)) + val first = new DummySparkPlanWithBatchScanChild( + outputPartitioning = KeyedPartitioning(Seq(id, t1), keys)) + val second = new DummySparkPlanWithBatchScanChild( + outputPartitioning = KeyedPartitioning(Seq(other, id, t1), paddedKeys)) + val distribution = ClusteredDistribution(Seq(id, t1)) + + withSQLConf( + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { + // Not a join, so `checkKeyGroupCompatible` declines and both children go through the + // per-child branch that pushes the positions down. + val parent = DummySparkPlan( + children = Seq(first, second), + requiredChildDistribution = Seq(distribution, distribution), + requiredChildOrdering = Seq(Nil, Nil)) + val planned = EnsureRequirements.apply(parent) + + assert(groupPartitionsNodes(planned.children.head).map(_.joinKeyPositions) === + Seq(Some(Seq(0, 1)))) + assert(groupPartitionsNodes(planned.children(1)).map(_.joinKeyPositions) === + Seq(Some(Seq(1, 2))), + "the positions pushed into a child must index into that child's own expressions") + } + } + test("SPARK-58968: a grouped KeyedPartitioning must still honour requiredNumPartitions") { val exprKey = AttributeReference("k", IntegerType)() // A grouped KeyedPartitioning with three distinct keys and three partitions. From f6a144284f43eb7c402afc75787a790d86a8e45b Mon Sep 17 00:00:00 2001 From: Peter Toth Date: Sat, 5 Sep 2026 08:49:11 +0200 Subject: [PATCH 2/3] [SPARK-59080][SQL] Address review: the cogroup path is reachable, not latent @sunchao pointed out that Pandas and Arrow cogroups group on real columns, unlike the Scala `CoGroupExec`, whose key comes from an `AppendColumns` that no `KeyedPartitioning` satisfies. Two keyed children do reach the per-child branch, so pushing each child's own `joinKeyPositions` is a reachable correctness fix rather than a latent one. The plan test is rebuilt on `FlatMapCoGroupsInPandasExec` instead of a synthetic parent. Both sides declare the same two key columns in the opposite order, so the cogroup key sits at position 1 on the left and at position 0 on the right, and both project onto the same key set. Without the production change the right side is handed the left side's positions and ends up grouped on its other partition column: `List(Some(List(1)), Some(List(1)))` where `List(Some(List(1)), Some(List(0)))` is right. --- .../exchange/EnsureRequirementsSuite.scala | 66 +++++++++++-------- 1 file changed, 37 insertions(+), 29 deletions(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala index d29b0a8797c24..42011ff00ccc1 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala @@ -1462,40 +1462,48 @@ class EnsureRequirementsSuite extends SharedSparkSession { ) test("SPARK-59080: pushed-down positions index into the child's own partition expressions") { - val id = AttributeReference("id", IntegerType)() - val t1 = AttributeReference("t1", IntegerType)() - val other = AttributeReference("other", IntegerType)() - // Both children carry the same (id, t1) key set, but the second declares a leading partition - // expression the distribution does not cluster on. So the keys it projects onto sit at - // positions 1 and 2, while the first child's sit at 0 and 1. - // - // No query reaches two keyed children here: a join is handled by `checkKeyGroupCompatible`, - // which pushes each side's own positions, and a cogroup's grouping key is synthesized so - // neither side stays keyed. The test pins the invariant rather than reproducing a query. - val keys = Seq(InternalRow(1, 1), InternalRow(1, 2), InternalRow(2, 1)) - val paddedKeys = Seq(InternalRow(9, 1, 1), InternalRow(9, 1, 2), InternalRow(9, 2, 1)) - val first = new DummySparkPlanWithBatchScanChild( - outputPartitioning = KeyedPartitioning(Seq(id, t1), keys)) - val second = new DummySparkPlanWithBatchScanChild( - outputPartitioning = KeyedPartitioning(Seq(other, id, t1), paddedKeys)) - val distribution = ClusteredDistribution(Seq(id, t1)) + val nL = AttributeReference("nL", IntegerType)() + val iL = AttributeReference("iL", IntegerType)() + val iR = AttributeReference("iR", IntegerType)() + val nR = AttributeReference("nR", IntegerType)() + // Both sides declare the same two key columns, in the opposite order, so the cogroup key sits + // at position 1 on the left and at position 0 on the right. Both project onto {1, 2}, so the + // sides are co-partitioned and neither is re-shuffled. + val leftKeys = Seq(InternalRow(1, 1), InternalRow(2, 1), InternalRow(3, 2)) + val rightKeys = Seq(InternalRow(1, 1), InternalRow(1, 2), InternalRow(2, 3)) + val left = new DummySparkPlanWithBatchScanChild( + outputPartitioning = KeyedPartitioning(Seq(nL, iL), leftKeys)) + val right = new DummySparkPlanWithBatchScanChild( + outputPartitioning = KeyedPartitioning(Seq(iR, nR), rightKeys)) + + val pythonUdf = PythonUDF("pyUDF", null, + StructType(Seq(StructField("value", IntegerType))), + Seq.empty, + PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF, + true) + // A cogroup requires `ClusteredDistribution` on both children but is not a `ShuffledJoin`, so + // `checkKeyGroupCompatible` declines and both children go through the per-child branch that + // pushes the positions down. Unlike the Scala `CoGroupExec`, whose key comes from an + // `AppendColumns` no `KeyedPartitioning` satisfies, the Pandas one groups on real columns, so + // both sides stay keyed. + val cogroup = FlatMapCoGroupsInPandasExec( + Seq(iL), Seq(iR), pythonUdf, + AttributeReference("value", IntegerType)() :: Nil, left, right) withSQLConf( SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { - // Not a join, so `checkKeyGroupCompatible` declines and both children go through the - // per-child branch that pushes the positions down. - val parent = DummySparkPlan( - children = Seq(first, second), - requiredChildDistribution = Seq(distribution, distribution), - requiredChildOrdering = Seq(Nil, Nil)) - val planned = EnsureRequirements.apply(parent) - - assert(groupPartitionsNodes(planned.children.head).map(_.joinKeyPositions) === - Seq(Some(Seq(0, 1)))) - assert(groupPartitionsNodes(planned.children(1)).map(_.joinKeyPositions) === - Seq(Some(Seq(1, 2))), + val result = EnsureRequirements.apply(cogroup) + + assert(result.collect { case s: ShuffleExchangeExec => s }.isEmpty, + "the sides are co-partitioned on the cogroup key") + assert(groupPartitionsNodes(result).map(_.joinKeyPositions) === + Seq(Some(Seq(1)), Some(Seq(0))), "the positions pushed into a child must index into that child's own expressions") + assert(result.children.map(_.outputPartitioning).forall { + case k: KeyedPartitioning => k.expressions == Seq(iL) || k.expressions == Seq(iR) + case _ => false + }, "each side must end up grouped on its own cogroup key") } } From 4799ca4256b6ceef73a665b3b3be35ba58f2c499 Mon Sep 17 00:00:00 2001 From: Peter Toth Date: Fri, 4 Sep 2026 19:21:20 +0200 Subject: [PATCH 3/3] [SPARK-59256][SQL] Choose a ShuffleSpecCollection member by pairing, not by enumeration order A `PartitioningCollection` offers several layouts, and which one is right depends on what the other side matched. Four places decided it by enumeration order instead, and none of them could see both sides. SPARK-59080 fixed the first, the per-child branch that shuffles a child. This addresses the other three, and makes the type say why a collection cannot answer alone. **The type split.** A new `LeafShuffleSpec` sub-trait carries `numPartitions` and `createPartitioning`, the seven concrete specs extend it, and `ShuffleSpecCollection` extends `ShuffleSpec` alone. `flatten` moves from a private helper in `EnsureRequirements` onto the hierarchy and returns `Seq[LeafShuffleSpec]`. `ShuffleSpec` is sealed, which is not what makes `flatten` total - it is a virtual method with an implementation per kind. What sealing buys is that "every spec is one layout or a choice of layouts" becomes a guarantee, which is what makes that return type honest, and that `SinglePartitionShuffleSpec`'s two-case match compiles as exhaustive. The collection's `createPartitioning` had no production caller after SPARK-59080, and only a runtime `require` stood between a future caller and a wrong partitioning. `numPartitions` had two consumers wanting two different aggregations, which is why the answer belongs to each caller: - `EnsureRequirements` ranks the children by parallelism to pick the reference layout. It wants the max, and it now takes it over the flattened members. - `SinglePartitionShuffleSpec.isCompatibleWith` asks whether the other side is a single partition, and it now answers `forall` over the members rather than reading the head's count. Not the `exists` the other specs use for a collection: they ask whether *some* member matches them and then plan on that member, while this asks a property of the child itself. The child is whichever member the other side picks, so a single partition has to hold for every one of them. Reading the head could claim a co-partitioning that does not hold, because a member's count is the count *after* the projection `createShuffleSpec` promises, not the child's own. **The pairing.** `createKeyedShuffleSpec` collapsed a collection to one member with `collectFirst`, per side, before either side had seen the other. Two sides can then pick members that do not agree, `checkKeyGroupCompatible` declines, and the join loses the storage-partitioned pushdown even though a pairing existed. It now returns every member's spec, and `checkKeyGroupCompatible` picks the pair that agrees on the keys and offers the most parallelism. The pick cannot be an independent per-side finest: a side whose only members are coarse would then fail to pair. It is also one rule and not two - a pair that is co-partitioned as it stands is not preferred over a finer one that needs a grouping node. The zero-node path is not what that would protect: it needs `left.outputPartitioning.exists(_ == leftPartitioning)`, which only an identity projection satisfies, and an identity projection reports the side's own physical count, which is maximal. `maxByOption` keeps the earliest element on a tie, and the earliest pair is the old per-side pick, so the pairing never loses the as-is path that pick used to find. What the ranking does trade, measured rather than hypothetical: between two pairs that both need grouping nodes, it takes 4 partitions with 2 empty ones on one side over 2 exactly-matched partitions. More parallelism for some padding. Everything after that derivation works on the chosen pair as before. One path changes reachability rather than behaviour: `reducersBothWays` now runs on a pair the old code could not form, so a connector whose `Reducer.resultType()` violates the `r(f1(x)) = f2(x)` contract can raise `storagePartitionJoinIncompatibleReducedTypesError` where the join used to fall back silently. That is what the error exists to catch, but it is a new failure mode. When no pair agrees, each side's first member is reported, which is what the per-side pick took, and the checks below fail on it exactly as they did before. **What changes at a default configuration, which is more than I first thought.** `pushPartValues` is on by default, and it is all the push branch in `checkKeyGroupCompatible` needs, so **the pairing changes plans out of the box**. Measured on the transform-difference shape with no config set at all: the per-side pick declines, `bestSpecOpt` is then empty because a keyed spec cannot be a shuffle reference without `v2BucketingShuffleEnabled`, which *is* off by default, and both children are shuffled onto the default 200 partitions. The pairing finds the agreeing pair and the join runs with no shuffle and two grouping nodes. The other two changes are no-ops unless `allowKeysSubsetOfPartitionKeys` is on: `PartitioningCollection` requires its members to agree on `numPartitions`, and every spec reports its own partitioning's count, so `max` equals `head` and the `forall` equals the head read. For the pairing to matter at all, a collection's members have to differ in a way `areKeysCompatible` can see, and there are two ways. `requireAllClusterKeysForCoPartition` off lets members cover different clustering keys. A transform difference does it with that requirement at its default, because a collection forces its members to agree on the key rows, and through those on the transforms' result types, but not on the transforms themselves. **Smaller things.** `CoalescedHashShuffleSpec.from` is typed `LeafShuffleSpec`, since it is structurally the one spec the coalesced spec was built from. Refining the declared return types of `HashPartitioning`, `NullAwareHashPartitioning` and `KeyedPartitioning`'s `createShuffleSpec` to their own spec types is what lets that be a type rather than a cast, and the `KeyedPartitioning` one deletes a production cast in `createKeyedShuffleSpecs` and two more in tests. `KeyedShuffleSpec.isCompatibleWith` read `other.numPartitions` inside a branch that had already matched `other` as a `KeyedShuffleSpec`, and now reads the narrowed value, the same object. `ShuffleSpecCollection` takes over the `require(specs.nonEmpty, ...)` that `numPartitions` used to carry, which is what keeps `flatten` non-empty and the ranking's `max` total. A dangling pre-split scaladoc for `ShuffleSpec`, stranded above `ShufflePartitionIdPassThrough`, is deleted rather than left to contradict the new trait doc. Planning cost: the per-side derivation is now linear in the number of satisfying members instead of stopping at the first, and the pairing is their cross product. Bounded in practice, since members of one collection reference different attributes and usually only one satisfies a given join's clustering. --- .../plans/physical/partitioning.scala | 104 +++++----- .../spark/sql/catalyst/ShuffleSpecSuite.scala | 65 +++++-- .../datasources/v2/GroupPartitionsExec.scala | 7 +- .../exchange/EnsureRequirements.scala | 75 ++++--- .../KeyGroupedPartitioningSuite.scala | 4 +- .../v2/GroupPartitionsExecSuite.scala | 3 +- .../exchange/EnsureRequirementsSuite.scala | 184 ++++++++++++++++++ 7 files changed, 341 insertions(+), 101 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 2354cf69205a0..979643ec4825c 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 @@ -318,7 +318,7 @@ trait HashPartitioningLike extends Expression with Partitioning with Unevaluable case class HashPartitioning(expressions: Seq[Expression], numPartitions: Int) extends HashPartitioningLike { - override def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpec = + override def createShuffleSpec(distribution: ClusteredDistribution): HashShuffleSpec = HashShuffleSpec(this, distribution) /** @@ -364,7 +364,7 @@ case class NullAwareHashPartitioning(expressions: Seq[Expression], numPartitions } } - override def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpec = + override def createShuffleSpec(distribution: ClusteredDistribution): NullAwareHashShuffleSpec = NullAwareHashShuffleSpec(this, distribution) override protected def withNewChildrenInternal( @@ -791,7 +791,7 @@ case class KeyedPartitioning( if (isGrouped) keysSatisfy(required) else mayGroupToSatisfy(required) } - override def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpec = { + override def createShuffleSpec(distribution: ClusteredDistribution): KeyedShuffleSpec = { val result = KeyedShuffleSpec(this, distribution) if (SQLConf.get.v2BucketingAllowKeysSubsetOfPartitionKeys) { // If allowing operation keys to be a subset of partition keys, create a new @@ -1198,18 +1198,6 @@ case class BroadcastPartitioning(mode: BroadcastMode) extends Partitioning { } } -/** - * This is used in the scenario where an operator has multiple children (e.g., join) and one or more - * of which have their own requirement regarding whether its data can be considered as - * co-partitioned from others. This offers APIs for: - * - * - Comparing with specs from other children of the operator and check if they are compatible. - * When two specs are compatible, we can say their data are co-partitioned, and Spark will - * potentially be able to eliminate shuffle if necessary. - * - Creating a partitioning that can be used to re-partition another child, so that to make it - * having a compatible partitioning as this node. - */ - /** * Represents a partitioning where partition IDs are passed through directly from the * DirectShufflePartitionID expression. This partitioning scheme is used when users @@ -1253,12 +1241,16 @@ case class ShufflePartitionIdPassThrough( copy(expr = newChildren.head.asInstanceOf[DirectShufflePartitionID]) } -trait ShuffleSpec { - /** - * Returns the number of partitions of this shuffle spec - */ - def numPartitions: Int - +/** + * Describes how a child's data is laid out, for the purpose of deciding whether two children are + * co-partitioned and, if not, what to shuffle the other one onto. + * + * A [[LeafShuffleSpec]] is one concrete layout. A [[ShuffleSpecCollection]] stands for a choice + * between several. A collection can answer [[isCompatibleWith]], which succeeds when any member + * matches. It cannot answer anything that needs one member: which one is right depends on what the + * other side matched, and only the caller comparing the two sides can see that. + */ +sealed trait ShuffleSpec { /** * Returns true iff this spec is compatible with the provided shuffle spec. * @@ -1271,10 +1263,28 @@ trait ShuffleSpec { def isCompatibleWith(other: ShuffleSpec): Boolean /** - * Whether this shuffle spec can be used to create partitionings for the other children. + * Whether this shuffle spec can be used to create partitionings for the other children. A + * [[ShuffleSpecCollection]] answers for the whole choice, since the planner asks it of a child's + * spec as a whole. Building the partitioning is [[LeafShuffleSpec.createPartitioning]], and that + * is always one member's job. */ def canCreatePartitioning: Boolean + /** + * This spec's leaf specs: a [[ShuffleSpecCollection]] yields its members recursively, and a + * [[LeafShuffleSpec]] yields itself. A caller that needs one member picks from these. Never + * empty, since a collection has at least one member. + */ + def flatten: Seq[LeafShuffleSpec] +} + +/** A [[ShuffleSpec]] describing one layout, as opposed to a choice between several. */ +trait LeafShuffleSpec extends ShuffleSpec { + /** + * Returns the number of partitions of this shuffle spec + */ + def numPartitions: Int + /** * Creates a partitioning that can be used to re-partition the other side with the given * clustering expressions. @@ -1284,11 +1294,19 @@ trait ShuffleSpec { */ def createPartitioning(clustering: Seq[Expression]): Partitioning = throw SparkUnsupportedOperationException() + + override final def flatten: Seq[LeafShuffleSpec] = this +: Nil } -case object SinglePartitionShuffleSpec extends ShuffleSpec { - override def isCompatibleWith(other: ShuffleSpec): Boolean = { - other.numPartitions == 1 +case object SinglePartitionShuffleSpec extends LeafShuffleSpec { + override def isCompatibleWith(other: ShuffleSpec): Boolean = other match { + case leaf: LeafShuffleSpec => leaf.numPartitions == 1 + // `forall`, not the `exists` the other specs use for a collection. They ask whether *some* + // member matches them, and the caller then plans on that member. This asks a property of the + // child itself, and the child is whichever member the other side picks, so a single partition + // has to be the answer for every one of them. The members can only disagree when the subset + // config projects them onto different key sets. + case ShuffleSpecCollection(specs) => specs.forall(isCompatibleWith) } override def canCreatePartitioning: Boolean = false @@ -1301,7 +1319,7 @@ case object SinglePartitionShuffleSpec extends ShuffleSpec { case class RangeShuffleSpec( numPartitions: Int, - distribution: ClusteredDistribution) extends ShuffleSpec { + distribution: ClusteredDistribution) extends LeafShuffleSpec { // `RangePartitioning` is not compatible with any other partitioning since it can't guarantee // data are co-partitioned for all the children, as range boundaries are randomly sampled. We @@ -1338,7 +1356,7 @@ private object HashShuffleSpecCompatibility { case class HashShuffleSpec( partitioning: HashPartitioning, - distribution: ClusteredDistribution) extends ShuffleSpec { + distribution: ClusteredDistribution) extends LeafShuffleSpec { /** * A sequence where each element is a set of positions of the hash partition key to the cluster @@ -1424,7 +1442,7 @@ case class HashShuffleSpec( */ case class NullAwareHashShuffleSpec( partitioning: NullAwareHashPartitioning, - distribution: ClusteredDistribution) extends ShuffleSpec { + distribution: ClusteredDistribution) extends LeafShuffleSpec { lazy val hashKeyPositions: Seq[mutable.BitSet] = { val distKeyToPos = mutable.Map.empty[Expression, mutable.BitSet] @@ -1481,8 +1499,8 @@ case class NullAwareHashShuffleSpec( } case class CoalescedHashShuffleSpec( - from: ShuffleSpec, - partitions: Seq[CoalescedBoundary]) extends ShuffleSpec { + from: LeafShuffleSpec, + partitions: Seq[CoalescedBoundary]) extends LeafShuffleSpec { override def isCompatibleWith(other: ShuffleSpec): Boolean = other match { case SinglePartitionShuffleSpec => @@ -1558,7 +1576,7 @@ case class IdentityReducer(transform: TransformExpression) extends Reducer[Any, case class KeyedShuffleSpec( partitioning: KeyedPartitioning, distribution: ClusteredDistribution, - joinKeyPositions: Option[Seq[Int]] = None) extends ShuffleSpec { + joinKeyPositions: Option[Seq[Int]] = None) extends LeafShuffleSpec { /** * A sequence where each element is a set of positions of the partition expression to the cluster @@ -1595,7 +1613,7 @@ case class KeyedShuffleSpec( // 4. the partition values from both sides are following the same order. case otherSpec @ KeyedShuffleSpec(otherPartitioning, otherDistribution, _) => distribution.clustering.length == otherDistribution.clustering.length && - numPartitions == other.numPartitions && areKeysCompatible(otherSpec) && + numPartitions == otherSpec.numPartitions && areKeysCompatible(otherSpec) && partitioning.partitionKeys == otherPartitioning.partitionKeys case ShuffleSpecCollection(specs) => specs.exists(isCompatibleWith) @@ -1775,7 +1793,7 @@ case class KeyedShuffleSpec( case class ShufflePartitionIdPassThroughSpec( partitioning: ShufflePartitionIdPassThrough, - distribution: ClusteredDistribution) extends ShuffleSpec { + distribution: ClusteredDistribution) extends LeafShuffleSpec { /** * A sequence where each element is a set of positions of the partition key to the cluster @@ -1818,7 +1836,14 @@ case class ShufflePartitionIdPassThroughSpec( override def numPartitions: Int = partitioning.numPartitions } +/** + * A choice between several layouts, produced by [[PartitioningCollection.createShuffleSpec]]. + * + * `specs` can hold a nested collection, since a [[PartitioningCollection]] can hold a nested one. + */ case class ShuffleSpecCollection(specs: Seq[ShuffleSpec]) extends ShuffleSpec { + require(specs.nonEmpty, "expected specs to be non-empty") + override def isCompatibleWith(other: ShuffleSpec): Boolean = { specs.exists(_.isCompatibleWith(other)) } @@ -1826,16 +1851,5 @@ case class ShuffleSpecCollection(specs: Seq[ShuffleSpec]) extends ShuffleSpec { override def canCreatePartitioning: Boolean = specs.forall(_.canCreatePartitioning) - override def createPartitioning(clustering: Seq[Expression]): Partitioning = { - // as we only consider # of partitions as the cost now, it doesn't matter which one we choose - // since they should all have the same # of partitions. - require(specs.map(_.numPartitions).toSet.size == 1, "expected all specs in the collection " + - "to have the same number of partitions") - specs.head.createPartitioning(clustering) - } - - override def numPartitions: Int = { - require(specs.nonEmpty, "expected specs to be non-empty") - specs.head.numPartitions - } + override def flatten: Seq[LeafShuffleSpec] = specs.flatMap(_.flatten) } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala index 496b408520fbc..4ab1d46d5fb35 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala @@ -59,7 +59,7 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper { } protected def checkCreatePartitioning( - spec: ShuffleSpec, + spec: LeafShuffleSpec, dist: ClusteredDistribution, expected: Partitioning): Unit = { val actual = spec.createPartitioning(dist.clustering) @@ -506,7 +506,6 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper { withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { val spec = reduced.createShuffleSpec(ClusteredDistribution(Seq(a))) - .asInstanceOf[KeyedShuffleSpec] assert(spec.joinKeyPositions === Some(Seq(0))) assert(spec.partitioning.partitionKeys.map(_.row.getInt(0)) === Seq(2020, 2021)) } @@ -613,13 +612,6 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper { SinglePartition ) - checkCreatePartitioning(ShuffleSpecCollection(Seq( - HashShuffleSpec(HashPartitioning(Seq($"a"), 10), distribution), - RangeShuffleSpec(10, distribution))), - ClusteredDistribution(Seq($"c", $"d")), - HashPartitioning(Seq($"c"), 10) - ) - // unsupported cases checkError( @@ -629,7 +621,7 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper { condition = "UNSUPPORTED_CALL.WITHOUT_SUGGESTION", parameters = Map( "methodName" -> "createPartitioning$", - "className" -> "org.apache.spark.sql.catalyst.plans.physical.ShuffleSpec")) + "className" -> "org.apache.spark.sql.catalyst.plans.physical.LeafShuffleSpec")) } test("compatibility: ShufflePartitionIdPassThroughSpec on both sides") { @@ -706,18 +698,57 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper { withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { val spec = collection.createShuffleSpec(ClusteredDistribution(Seq(id, t1))) - .asInstanceOf[ShuffleSpecCollection] // The disagreement is kept rather than resolved here. Every member has to stay for // `isCompatibleWith`, which answers for any of them, and the collection cannot know which one - // the other side matched. `EnsureRequirements` resolves that and asks the member, not the - // collection. - assert(spec.specs.map(_.numPartitions).toSet === Set(3, 2)) + // the other side matched. `EnsureRequirements` resolves that and asks the member, so the + // collection has no partition count of its own to read. + val memberPartitions = spec.flatten.map(_.numPartitions) + assert(memberPartitions.toSet === Set(3, 2)) assert(spec.isCompatibleWith(spec), "every member stays available for matching") + } + } - // So asking the collection for a single answer is the caller's mistake, and it says so. - val e = intercept[IllegalArgumentException](spec.createPartitioning(Seq(id, t1))) - assert(e.getMessage.contains("expected all specs in the collection to have the same number")) + test("SPARK-59256: a single-partition side needs every member of a collection to be one") { + val id = AttributeReference("id", IntegerType)() + val t1 = AttributeReference("t1", IntegerType)() + val t2 = AttributeReference("t2", IntegerType)() + // Clustering on (id, t1) again. The first member matches `id` only and collapses to one + // partition, the second projects onto both positions and keeps three. The child itself has + // three, so it is not co-partitioned with a single partition, whatever the head says. + val keys = Seq(InternalRow(1, 1), InternalRow(1, 2), InternalRow(1, 3)) + val collection = PartitioningCollection.fromPartitionings(Seq( + KeyedPartitioning(Seq(id, t2), keys), + KeyedPartitioning(Seq(id, t1), keys))) + + withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { + val spec = collection.createShuffleSpec(ClusteredDistribution(Seq(id, t1))) + val memberPartitions = spec.flatten.map(_.numPartitions) + // Ordered, not just as a set: the whole point is that the head is the wrong one to read. + assert(memberPartitions === Seq(1, 3)) + + // Reading the head's count answers one partition and claims a co-partitioning that does not + // hold. Only this direction is asserted: `KeyedShuffleSpec` has no + // `SinglePartitionShuffleSpec` case at all, so the reverse is false whatever the member + // counts are. That asymmetry is pre-existing and is not what this change is about. + assert(!SinglePartitionShuffleSpec.isCompatibleWith(spec)) } } + + test("SPARK-59256: an empty collection is rejected at construction, not at the first read") { + // `numPartitions` carried this `require` and threw the same way, but only once something asked. + // `flatten` and the ranking's `max` both rely on a collection being non-empty. + val e = intercept[IllegalArgumentException](ShuffleSpecCollection(Nil)) + assert(e.getMessage.contains("expected specs to be non-empty")) + } + + test("SPARK-59256: flattening reaches the members of a nested collection") { + val distribution = ClusteredDistribution(Seq($"a", $"b")) + val buried = HashShuffleSpec(HashPartitioning(Seq($"a"), 10), distribution) + val direct = HashShuffleSpec(HashPartitioning(Seq($"b"), 10), distribution) + // A `PartitioningCollection` can hold another one, so a spec collection can nest too. + val collection = ShuffleSpecCollection(Seq(ShuffleSpecCollection(Seq(buried)), direct)) + + assert(collection.flatten === Seq(buried, direct)) + } } 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..36669d84cf7a6 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 @@ -87,9 +87,10 @@ case class GroupPartitionsExec( assert(projectedExpressions.length == exprs.length) projectedExpressions.zip(exprs).map { case (expr, Some(KeyReducer(_, reduced))) => - // `reduced` was stored from the single spec that `createKeyedShuffleSpec` - // picked (`collectFirst`); re-target it at this `KeyedPartitioning`'s own key - // attribute so that every `KeyedPartitioning` in a collection keeps its own. + // `reduced` came from the one member `checkKeyGroupCompatible` paired this + // side on, which need not be the member being rewritten. The keys are reduced + // once, from the shared key rows, so `reduced` describes them whichever member + // this is, and only the key attribute has to be re-targeted. reduced.withReference(expr.references.head) case (expr, None) => expr } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala index 65f79ddac099a..ae9f3f3a8720e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala @@ -218,8 +218,9 @@ case class EnsureRequirements( } else { candidateSpecs } - // Pick the spec with the best parallelism - Some(finalCandidateSpecs.values.maxBy(_.numPartitions)) + // Pick the spec with the best parallelism. For a collection that is the best any member + // offers, since reading one member's count would depend on the enumeration order. + Some(finalCandidateSpecs.values.maxBy(_.flatten.map(_.numPartitions).max)) } // Check if the following conditions are satisfied: @@ -256,11 +257,11 @@ case class EnsureRequirements( childrenIndexes.filter(i => best.isCompatibleWith(specs(i))) } lazy val bestMemberOpt = bestSpecOpt.flatMap { best => - val matchedMembers = matchedIndexes.map(i => flattenSpec(specs(i))) + val matchedMembers = matchedIndexes.map(i => specs(i).flatten) // No member serving every matched child means there is no layout to align them on, so they // all take the ordinary shuffle. That needs three or more clustered children, since with // two the member that reported the match serves both, and no operator has three today. - flattenSpec(best) + best.flatten .filter(m => matchedMembers.forall(_.exists(m.isCompatibleWith))) .maxByOption(_.numPartitions) } @@ -275,7 +276,7 @@ case class EnsureRequirements( // own partition expressions -- the chosen best member only says which member of it the // two sides agreed on. val bestMember = bestMemberOpt.get - flattenSpec(specs(idx)).find(bestMember.isCompatibleWith) match { + specs(idx).flatten.find(bestMember.isCompatibleWith) match { // If `areChildrenCompatible` is false, we can still perform SPJ // by shuffling the other side based on join keys (see the else case below). // Hence we need to ensure that after this call, the outputPartitioning of the @@ -520,16 +521,30 @@ case class EnsureRequirements( var newLeft = left var newRight = right - val specs = Seq(left, right).zip(requiredChildDistribution).map { case (p, d) => - if (!d.isInstanceOf[ClusteredDistribution]) return None - val cd = d.asInstanceOf[ClusteredDistribution] - val specOpt = createKeyedShuffleSpec(p.outputPartitioning, cd) - if (specOpt.isEmpty) return None - specOpt.get - } - - val leftSpec = specs.head - val rightSpec = specs(1) + def candidatesFor(plan: SparkPlan, required: Distribution): Seq[KeyedShuffleSpec] = + required match { + case cd: ClusteredDistribution => createKeyedShuffleSpecs(plan.outputPartitioning, cd) + case _ => Nil + } + val leftCandidates = candidatesFor(left, requiredChildDistribution.head) + val rightCandidates = candidatesFor(right, requiredChildDistribution(1)) + if (leftCandidates.isEmpty || rightCandidates.isEmpty) return None + + // Each side may offer several members, and the right one is the one the other side can pair + // with, which neither side can tell on its own. So pick the pair rather than a member per side, + // and rank the pairs that agree on the keys by the parallelism they offer, the same trade + // `ensureDistributionAndOrdering` makes between children when it picks `bestSpecOpt`. + val agreeingPairs = for { + l <- leftCandidates + r <- rightCandidates + if l.areKeysCompatible(r) + } yield (l, r) + val (leftSpec, rightSpec) = agreeingPairs + .maxByOption { case (l, r) => l.numPartitions.max(r.numPartitions) } + // No agreeing pair means every pair fails the checks below, so the method returns `None` + // whichever one it reports. Reporting each side's first member keeps that path byte for byte + // what the per-side pick produced, `logInfo` included. + .getOrElse((leftCandidates.head, rightCandidates.head)) val leftPartitioning = leftSpec.partitioning val rightPartitioning = rightSpec.partitioning @@ -537,7 +552,7 @@ case class EnsureRequirements( // partitionings are not modified (projected) in specs and left and right side partitionings are // compatible with each other. // Left and right `outputPartitioning` is a `PartitioningCollection` or a `KeyedPartitioning` - // otherwise `createKeyedShuffleSpec()` would have returned `None`. + // otherwise `createKeyedShuffleSpecs()` would have returned nothing. var isCompatible = left.outputPartitioning.asInstanceOf[Expression].exists(_ == leftPartitioning) && right.outputPartitioning.asInstanceOf[Expression].exists(_ == rightPartitioning) && @@ -702,7 +717,7 @@ case class EnsureRequirements( val originalPartitioning = partiallyClusteredChild.outputPartitioning.asInstanceOf[Expression] // `outputPartitioning` is either a `PartitioningCollection` or a `KeyedPartitioning` - // otherwise `createKeyedShuffleSpec()` would have returned `None`. + // otherwise `createKeyedShuffleSpecs()` would have returned nothing. val originalKeyedPartitioning = originalPartitioning.collectFirst { case k: KeyedPartitioning => k }.get val projectedOriginalPartitionKeys = partiallyClusteredSpec.joinKeyPositions @@ -798,12 +813,6 @@ case class EnsureRequirements( } } - // Flattens a (possibly nested) `ShuffleSpecCollection` into its member specs. - private def flattenSpec(spec: ShuffleSpec): Seq[ShuffleSpec] = spec match { - case ShuffleSpecCollection(specs) => specs.flatMap(flattenSpec) - case other => Seq(other) - } - /** * Applies join key positions to a plan by wrapping or updating GroupPartitionsExec. */ @@ -818,13 +827,15 @@ case class EnsureRequirements( } /** - * Tries to create a [[KeyedShuffleSpec]] from the input partitioning and distribution, if the - * partitioning is a [[KeyedPartitioning]] (either directly or indirectly), and satisfies the - * given distribution. + * Every [[KeyedShuffleSpec]] the input partitioning can offer for the given distribution, one per + * [[KeyedPartitioning]] in it that satisfies it. A [[PartitioningCollection]] yields them in + * member order, nested collections included, and the caller picks. Returning only the first would + * decide by enumeration order which member the join is planned on, and only the caller comparing + * the two sides knows which member pairs with the other side's. */ - private def createKeyedShuffleSpec( + private def createKeyedShuffleSpecs( partitioning: Partitioning, - distribution: ClusteredDistribution): Option[KeyedShuffleSpec] = { + distribution: ClusteredDistribution): Seq[KeyedShuffleSpec] = { def tryCreate(partitioning: KeyedPartitioning): Option[KeyedShuffleSpec] = { // The config requires all the cluster keys to be covered by the partition keys, to avoid // the skew of joining on keys that are coarser than the join keys. Key order and duplicated @@ -839,17 +850,17 @@ case class EnsureRequirements( if (partitioning.satisfies(distribution) && (!SQLConf.get.getConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION) || allClusterKeysCovered)) { - Some(partitioning.createShuffleSpec(distribution).asInstanceOf[KeyedShuffleSpec]) + Some(partitioning.createShuffleSpec(distribution)) } else { None } } partitioning match { - case p: KeyedPartitioning => tryCreate(p) + case p: KeyedPartitioning => tryCreate(p).toSeq case PartitioningCollection(partitionings) => - partitionings.collectFirst(Function.unlift(createKeyedShuffleSpec(_, distribution))) - case _ => None + partitionings.flatMap(createKeyedShuffleSpecs(_, distribution)) + case _ => Nil } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala index c700ce6cbd3d3..858744198ac88 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala @@ -887,8 +887,8 @@ class KeyGroupedPartitioningSuite test("SPARK-59045: reduced expression is retargeted per KeyedPartitioning") { // A chained SPJ's output partitioning reports one `KeyedPartitioning` per join side, but the - // reduced expression is derived from the single spec that `createKeyedShuffleSpec` picks - // (`collectFirst`). Re-targeting it at each `KeyedPartitioning`'s own key attribute keeps the + // reduced expression is derived from the one member `checkKeyGroupCompatible` paired this side + // on. Re-targeting it at each `KeyedPartitioning`'s own key attribute keeps the // other sides' partitionings intact - otherwise a GROUP BY on the other side's key no longer // sees a partitioning on it and the query shuffles (0 shuffles on base and here, 1 if the // use-site re-targeting is dropped). diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala index d8ea11c6ea37e..a023d85d71ce5 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala @@ -20,7 +20,7 @@ package org.apache.spark.sql.execution.datasources.v2 import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeReference, SortOrder, TransformExpression} -import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, KeyedPartitioning, KeyedShuffleSpec, KeyReducer, Partitioning, PartitioningCollection, UnknownPartitioning} +import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, KeyedPartitioning, KeyReducer, Partitioning, PartitioningCollection, UnknownPartitioning} import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper import org.apache.spark.sql.connector.catalog.functions.{BucketFunction, BucketReducer} import org.apache.spark.sql.execution.{DummySparkPlan, LeafExecNode, SafeForKWayMerge} @@ -334,7 +334,6 @@ class GroupPartitionsExecSuite extends SharedSparkSession { withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { val spec = partitioning.createShuffleSpec(ClusteredDistribution(Seq(exprA))) - .asInstanceOf[KeyedShuffleSpec] assert(spec.joinKeyPositions === Some(Seq(0))) val gpe = GroupPartitionsExec( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala index 42011ff00ccc1..ae239839a9dc0 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala @@ -1507,6 +1507,190 @@ class EnsureRequirementsSuite extends SharedSparkSession { } } + test("SPARK-59256: a collection is ranked on its best member, not on whichever came first") { + val id = AttributeReference("id", IntegerType)() + val t1 = AttributeReference("t1", IntegerType)() + val t2 = AttributeReference("t2", IntegerType)() + // Clustering on (id, t1). The collection's first member matches `id` only and keeps two + // partitions, the second covers both positions and keeps five. Reading the first member's + // count ranks this child at 2 and loses to the hashed child's 3, even though the member the + // shuffle would actually be built from offers 5. + val keys = Seq(InternalRow(1, 1), InternalRow(1, 2), InternalRow(1, 3), + InternalRow(2, 1), InternalRow(2, 2)) + val keyed = new DummySparkPlanWithBatchScanChild( + outputPartitioning = PartitioningCollection.fromPartitionings(Seq( + KeyedPartitioning(Seq(id, t2), keys), + KeyedPartitioning(Seq(id, t1), keys)))) + val hashed = DummySparkPlan(outputPartitioning = HashPartitioning(Seq(id, t1), 3)) + val distribution = ClusteredDistribution(Seq(id, t1)) + + withSQLConf( + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { + // Pin the premise, ordered: the coarse member really is the one a head read would take. + val memberPartitions = keyed.outputPartitioning.createShuffleSpec(distribution) + .flatten.map(_.numPartitions) + assert(memberPartitions === Seq(2, 5)) + + val parent = DummySparkPlan( + children = Seq(keyed, hashed), + requiredChildDistribution = Seq(distribution, distribution), + requiredChildOrdering = Seq(Nil, Nil)) + val planned = EnsureRequirements.apply(parent) + assert(planned.children.head.collect { case s: ShuffleExchangeExec => s }.isEmpty, + "the keyed side wins the ranking, so it is not re-shuffled") + val shuffles = planned.children(1).collect { case s: ShuffleExchangeExec => s } + assert(shuffles.map(_.outputPartitioning.numPartitions) === Seq(5), + "the hashed side lands on the keyed side's finest member") + } + } + + test("SPARK-59256: the join is planned on a member pair, not on each side's first member") { + val aL = AttributeReference("aL", IntegerType)() + val bL = AttributeReference("bL", IntegerType)() + val cL = AttributeReference("cL", IntegerType)() + val aR = AttributeReference("aR", IntegerType)() + val bR = AttributeReference("bR", IntegerType)() + val cR = AttributeReference("cR", IntegerType)() + // The join keys are wider than the partitioning arity, so a member can satisfy the distribution + // without covering every join key. + val keys = Seq(InternalRow(1, 1), InternalRow(1, 2), InternalRow(2, 1), InternalRow(2, 2)) + val left = new DummySparkPlanWithBatchScanChild( + outputPartitioning = PartitioningCollection.fromPartitionings(Seq( + KeyedPartitioning(Seq(aL, bL), keys), + KeyedPartitioning(Seq(aL, cL), keys)))) + // The b-member and the c-member are in the opposite order here. + val right = new DummySparkPlanWithBatchScanChild( + outputPartitioning = PartitioningCollection.fromPartitionings(Seq( + KeyedPartitioning(Seq(aR, cR), keys), + KeyedPartitioning(Seq(aR, bR), keys)))) + + withSQLConf( + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", + SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") { + val smj = SortMergeJoinExec( + Seq(aL, bL, cL), Seq(aR, bR, cR), Inner, None, left, right) + val planned = EnsureRequirements.apply(smj) + + assert(planned.collect { case s: ShuffleExchangeExec => s }.isEmpty, + "neither side needs a shuffle either way") + // Taking each side's first member pairs (aL, bL) with (aR, cR), which do not agree on their + // keys, so `checkKeyGroupCompatible` declines and each side is merely grouped on its own + // keys: two `GroupPartitionsExec` nodes, and the common partition values never pushed. + // Pairing finds (aL, bL) with (aR, bR), co-partitioned as they stand, so nothing is inserted. + assert(groupPartitionsNodes(planned).isEmpty, + "the paired members are already co-partitioned, so no grouping node is needed") + } + } + + test("SPARK-59256: the finest agreeing pair wins, not the first one") { + val aL = AttributeReference("aL", IntegerType)() + val bL = AttributeReference("bL", IntegerType)() + val zL = AttributeReference("zL", IntegerType)() + val aR = AttributeReference("aR", IntegerType)() + val bR = AttributeReference("bR", IntegerType)() + val wR = AttributeReference("wR", IntegerType)() + // Each side leads with a member whose second expression is not a join key, so it projects onto + // `a` alone, and follows with one covering both. Both pairings agree on the keys, so the choice + // is the ranking's: the coarse pair leaves 2 partitions a side, the fine one 4 and 2. + val leftKeys = + Seq(InternalRow(1, 10), InternalRow(1, 11), InternalRow(2, 20), InternalRow(2, 21)) + val rightKeys = Seq(InternalRow(1, 10), InternalRow(2, 20)) + val left = new DummySparkPlanWithBatchScanChild( + outputPartitioning = PartitioningCollection.fromPartitionings(Seq( + KeyedPartitioning(Seq(aL, zL), leftKeys), + KeyedPartitioning(Seq(aL, bL), leftKeys)))) + val right = new DummySparkPlanWithBatchScanChild( + outputPartitioning = PartitioningCollection.fromPartitionings(Seq( + KeyedPartitioning(Seq(aR, wR), rightKeys), + KeyedPartitioning(Seq(aR, bR), rightKeys)))) + + withSQLConf( + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", + SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") { + val smj = SortMergeJoinExec(Seq(aL, bL), Seq(aR, bR), Inner, None, left, right) + val planned = EnsureRequirements.apply(smj) + + assert(planned.collect { case s: ShuffleExchangeExec => s }.isEmpty, + "both pairings agree on the keys, so neither side is shuffled either way") + // Taking the first agreeing pair groups both sides on `a` alone, i.e. positions `Seq(0)`. + assert(groupPartitionsNodes(planned).map(_.joinKeyPositions) === + Seq(Some(Seq(0, 1)), Some(Seq(0, 1))), + "the join must be planned on the pair that keeps the most partitions") + } + } + + test("SPARK-59256: the pairing reaches a default configuration through a transform difference") { + val aL = AttributeReference("aL", IntegerType)() + val bL = AttributeReference("bL", IntegerType)() + val aR = AttributeReference("aR", IntegerType)() + val bR = AttributeReference("bR", IntegerType)() + // A collection does not require its members to agree on the transform at a position, only on + // arity and on the key rows - and, through those rows, on the transforms' result types, which + // is why the two here are `years` and `bucket` rather than `years` and `days`. So the members + // can differ with the co-partition requirement left at its default, and only the second one + // pairs with the other side. + val leftKeys = Seq(InternalRow(1, 10), InternalRow(2, 20)) + val rightKeys = Seq(InternalRow(1, 10), InternalRow(3, 30)) + val left = new DummySparkPlanWithBatchScanChild( + outputPartitioning = PartitioningCollection.fromPartitionings(Seq( + KeyedPartitioning(Seq(days(aL), years(bL)), leftKeys), + KeyedPartitioning(Seq(days(aL), bucket(4, bL)), leftKeys)))) + val right = new DummySparkPlanWithBatchScanChild( + outputPartitioning = KeyedPartitioning(Seq(days(aR), bucket(4, bR)), rightKeys)) + + // No `withSQLConf` on purpose. `pushPartValues` is on by default, which is all the push branch + // needs, so this is what a user gets out of the box. + val smj = SortMergeJoinExec(Seq(aL, bL), Seq(aR, bR), Inner, None, left, right) + val planned = EnsureRequirements.apply(smj) + + // Without the pairing, the left's first member is taken and `years(bL)` is matched against + // `bucket(4, bR)`, which is not the same function, so the join declines and both sides are + // shuffled onto the default partitioning. + assert(planned.collect { case s: ShuffleExchangeExec => s }.isEmpty, + "the second member pairs with the other side, so neither side is shuffled") + assert(groupPartitionsNodes(planned).map(_.expectedPartitionKeys.map(_.size)) === + Seq(Some(3), Some(3)), + "both sides are pushed the union of the two key sets") + } + + test("SPARK-59256: no agreeing pair leaves the join alone however many members each side has") { + val aL = AttributeReference("aL", IntegerType)() + val bL = AttributeReference("bL", IntegerType)() + val aR = AttributeReference("aR", IntegerType)() + val bR = AttributeReference("bR", IntegerType)() + // The two sides differ in arity, which `areKeysCompatible` refuses whatever the keys are, so + // none of the four pairs agrees and the fallback reports each side's first member. A negative + // pin: it holds with the per-side pick too, which is the point - the fallback must not start + // claiming a co-partitioning the old code did not claim. + val oneKeys = Seq(InternalRow(1), InternalRow(2)) + val twoKeys = Seq(InternalRow(1, 10), InternalRow(2, 20)) + val left = new DummySparkPlanWithBatchScanChild( + outputPartitioning = PartitioningCollection.fromPartitionings(Seq( + KeyedPartitioning(Seq(aL), oneKeys), + KeyedPartitioning(Seq(bL), oneKeys)))) + val right = new DummySparkPlanWithBatchScanChild( + outputPartitioning = PartitioningCollection.fromPartitionings(Seq( + KeyedPartitioning(Seq(aR, bR), twoKeys), + KeyedPartitioning(Seq(bR, aR), twoKeys)))) + + withSQLConf( + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") { + val smj = SortMergeJoinExec(Seq(aL, bL), Seq(aR, bR), Inner, None, left, right) + val planned = EnsureRequirements.apply(smj) + + // The shuffle is what pins this, not the absence of grouping nodes: a fallback that wrongly + // claimed compatibility would return the children untouched and skip the push branch, so + // there would be no grouping node either way. + assert(planned.collect { case s: ShuffleExchangeExec => s }.nonEmpty, + "no pair agrees, so the join is not planned as storage-partitioned") + assert(groupPartitionsNodes(planned).isEmpty, "and nothing is grouped") + } + } + test("SPARK-58968: a grouped KeyedPartitioning must still honour requiredNumPartitions") { val exprKey = AttributeReference("k", IntegerType)() // A grouped KeyedPartitioning with three distinct keys and three partitions.