Skip to content

[SPARK-59080][SQL][4.2] Pick one ShuffleSpecCollection member for the SPJ pushdown and the re-shuffle #58565

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
peter-toth wants to merge 1 commit into from
Closed

[SPARK-59080][SQL][4.2] Pick one ShuffleSpecCollection member for the SPJ pushdown and the re-shuffle #58565

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 31 additions & 1 deletion sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -591,4 +591,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(Seq(
KeyedPartitioning(Seq(id, t1), keys),
KeyedPartitioning(Seq(id, t2), keys)))

withSQLConf(SQLConf.V2_BUCKETING_ALLOW_JOIN_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"))
}
}
}
47 changes: 32 additions & 15 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -776,12 +795,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)
}

/**
Expand Down
51 changes: 51 additions & 0 deletions sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3162,6 +3162,57 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
}
}

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_JOIN_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"))
Expand Down
78 changes: 78 additions & 0 deletions ...core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1397,6 +1397,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(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_JOIN_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)),
Expand All @@ -1405,6 +1437,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_JOIN_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.
Expand Down