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..ef0649a81ecc0 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 @@ -1109,7 +1109,7 @@ case class PartitioningCollection(partitionings: Seq[Partitioning]) override def satisfies0(required: Distribution): Boolean = partitionings.exists(_.satisfies(required)) - override def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpec = { + override def createShuffleSpec(distribution: ClusteredDistribution): ShuffleSpecCollection = { val filtered = partitionings.filter(_.satisfies(distribution)) ShuffleSpecCollection(filtered.map(_.createShuffleSpec(distribution))) } @@ -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,17 @@ 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. + * + * There are two kinds, and the split is deliberate. A [[LeafShuffleSpec]] is one concrete layout. + * A [[ShuffleSpecCollection]] stands for a choice between several, so it can answer + * [[isCompatibleWith]], which succeeds when any member matches, but it cannot answer anything that + * needs one member: which member is the right one depends on what the other side matched, and that + * is only visible to the caller comparing the two sides. + */ +sealed trait ShuffleSpec { /** * Returns true iff this spec is compatible with the provided shuffle spec. * @@ -1271,10 +1264,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 eligible 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 +1295,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 +1320,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 +1357,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 +1443,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 +1500,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 +1577,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 +1614,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 +1794,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 +1837,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 +1852,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 b6cf5dec1f5f7..d70ac4388c865 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 @@ -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) @@ -613,13 +613,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 +622,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") { @@ -690,4 +683,67 @@ 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))) + + // 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, 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") + } + } + + 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: 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)) + assert(buried.flatten === Seq(buried), "a leaf flattens to itself") + } } 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..4bd585e992af4 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. Every member reduces + // to the same values, since they share the key rows and the types those rows + // were built with, so 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 2a33280861da0..d8ad44a0d9e2b 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 @@ -219,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: @@ -247,30 +247,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 => 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. + best.flatten + .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 + 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 // 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 @@ -501,16 +521,31 @@ 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) } + // With no agreeing pair, every pair fails the checks below and the method returns `None` + // whichever one it reports, since `isCompatibleWith` has `areKeysCompatible` as a conjunct. + // Report each side's first member, which is what the per-side pick took, so the `logInfo` + // below stays on the path that used to emit it. + .getOrElse((leftCandidates.head, rightCandidates.head)) val leftPartitioning = leftSpec.partitioning val rightPartitioning = rightSpec.partitioning @@ -518,7 +553,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) && @@ -683,7 +718,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 @@ -779,14 +814,6 @@ 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 - } - /** * Applies join key positions to a plan by wrapping or updating GroupPartitionsExec. */ @@ -801,13 +828,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 @@ -822,17 +851,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 6a911ce652615..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). @@ -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..1b74fe054a0e3 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,169 @@ 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-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. That is the only way two members of one collection can + // differ in coverage, since the collection requires them to have matching arity. + 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, so taking each side's first + // member pairs (aL, bL) with (aR, cR), which do not agree on their keys. + 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-58968: a grouped KeyedPartitioning must still honour requiredNumPartitions") { val exprKey = AttributeReference("k", IntegerType)() // A grouped KeyedPartitioning with three distinct keys and three partitions.