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..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 @@ -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,52 @@ class EnsureRequirementsSuite extends SharedSparkSession { requiredChildOrdering = Seq(Seq.empty) ) + test("SPARK-59080: pushed-down positions index into the child's own partition expressions") { + 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") { + 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") + } + } + 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.