diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala index 191de9bf5062f..bfa94914f4ae9 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala @@ -123,16 +123,6 @@ case class TransformExpression( } } - /** - * Re-targets this partition transform expression at `attr`. A partition transform expression - * has a single leaf attribute (`KeyedPartitioning.supportsExpressions`), so this replaces that - * attribute and keeps the rest of the expression: any field path above the leaf comes from this - * expression, not from the re-targeted key. Re-targeting is therefore only faithful when the - * source and the target key expressions have the same path shape. - */ - def withReference(attr: Attribute): TransformExpression = - transform { case _: AttributeReference => attr }.asInstanceOf[TransformExpression] - // Return a Reducer for a reducible function on another reducible function private def reducer( thisFunction: ReducibleFunction[_, _], 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 ccb8a4994ec19..4a9f8cdd2a2c6 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 @@ -23,6 +23,7 @@ import scala.collection.mutable import org.apache.spark.{SparkException, SparkUnsupportedOperationException} import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.plans.QueryPlan import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper import org.apache.spark.sql.connector.catalog.functions.Reducer import org.apache.spark.sql.internal.SQLConf @@ -793,10 +794,11 @@ case class KeyedPartitioning( // The operation keys may be a subset of the partition keys, so one partition expression // covering one of them is enough. Partitions can then still hold rows that share an // operation key. What makes the distribution true is the projection onto the covering - // positions, plus the grouping that follows it. Given the single reference per - // expression that `supportsExpressions` enforces, the test below is exactly - // `KeyedShuffleSpec.keyPositions.exists(_.nonEmpty)`, where `createShuffleSpec` takes - // its `joinKeyPositions` from. Consolidating the two is left to a follow-up. + // positions, plus the grouping that follows it. An expression over several columns + // maps to no position, so a layout with one is refused. For any other layout, the + // test below is exactly `KeyedShuffleSpec.keyPositions.exists(_.nonEmpty)`, where + // `createShuffleSpec` takes its `joinKeyPositions` from. Consolidating the two is left + // to a follow-up. requiredClustering.exists(x => attributes.exists(_.semanticEquals(x))) && expressions.forall(_.references.size == 1) } else { @@ -850,6 +852,16 @@ case class KeyedPartitioning( if (mayContainUnknownPartitionKeys && joinKeyPositions.length < expressions.length) { return result } + // No position covers a cluster key, so there is nothing to project onto. Projecting onto + // none would coalesce every partition into one, which is never the answer, and it is the + // same refusal `EnsureRequirements.resolveKeyedPartitioning` makes for its own caller. + // `keyPositions` maps an expression's single *reference* onto the clustering, so this is + // reachable for a clustering that names the partition expression itself, where + // `keysSatisfy` says yes and this finds nothing. That covers an expression over several + // columns too, e.g. `b + c`. Return the unprojected spec, as above. + if (joinKeyPositions.isEmpty && expressions.nonEmpty) { + return result + } // If allowing operation keys to be a subset of partition keys, create a new // `KeyedPartitioning` grouped on the operation keys, and use that as // the returned shuffle spec. @@ -1334,12 +1346,22 @@ trait ShuffleSpec { * co-partitioned with the `other`, and therefore no shuffle is required when joining the two * sides. * - * Note that Spark assumes this to be reflexive, symmetric and transitive. + * Note that Spark assumes this to be reflexive, symmetric and transitive. Some specs are not + * reflexive: + * - a [[RangeShuffleSpec]], whose boundaries are sampled; + * - a [[KeyedShuffleSpec]] with a partition expression that `keyPositions` maps to no cluster + * key, e.g. `years(ts)` clustered on `years(ts)`, or `bucket(4, b + c)`; + * - a [[KeyedShuffleSpec]] with a partition expression that is neither a column nor a transform, + * e.g. `b + 1` or `s.a`. + * None of them answers [[canCreatePartitioning]], which keeps the planner from picking one as + * the layout to shuffle the other children onto. */ 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]] can when any of its members can. Building the partitioning is always + * one member's job, so the planner builds it from a member that can. */ def canCreatePartitioning: Boolean @@ -1584,10 +1606,10 @@ case class CoalescedHashShuffleSpec( * keys, and this is this side's own expression marked with the pairing * that reduced it (`TransformExpression.reducedWith`). Stored as-is from * the expression pair that produced the reducer; the only structural - * consumer, `GroupPartitionsExec`, re-targets it at each reported - * `KeyedPartitioning`'s own key attribute in `outputPartitioning` and + * consumer, `GroupPartitionsExec`, rebuilds it over each reported + * `KeyedPartitioning`'s own argument in `outputPartitioning` and * normalizes it positionally in `doCanonicalize`, so the attribute it - * carries here is not load-bearing. + * carries here does not matter. */ case class KeyReducer(reducer: Reducer[_, _], reducedExpression: TransformExpression) @@ -1600,11 +1622,13 @@ case class KeyReducer(reducer: Reducer[_, _], reducedExpression: TransformExpres * It is a case class so that structurally identical reducers compare by value, and plans holding * them stay canonicalization-equal. * - * @param transform the partition transform, re-targeted at the identity side's key attribute + * @param transform the partition transform, rebuilt over the identity side's key attribute */ case class IdentityReducer(transform: TransformExpression) extends Reducer[Any, Any] { - // `transform` has a single leaf attribute (`KeyedPartitioning.supportsExpressions`), which is - // bound to ordinal 0 of the single-value row `reduce` evaluates it against. + // `transform` is over a single bare column, which is bound to ordinal 0 of the single-value row + // `reduce` evaluates it against. `KeyedShuffleSpec.isExpressionCompatible` admits only bare + // column arguments here, and `keyPositions` maps an expression to a position only when it has + // one reference. @transient private lazy val bound: Expression = BindReferences.bindReference(transform, AttributeSeq(transform.references.toSeq)) @@ -1633,8 +1657,19 @@ case class KeyedShuffleSpec( * keys. For instance, if cluster keys are [a, b, b] and partition expressions are * [bucket(4, a), years(b)], the result will be [(0), (1, 2)]. * - * Note that we only allow each partition expression to contain a single partition key. - * Therefore the mapping here is very similar to that from `HashShuffleSpec`. + * Only an expression with a single reference maps to a position, so the mapping here is very + * similar to that from `HashShuffleSpec`. A scan reports no other expression + * (`KeyedPartitioning.supportsExpressions`), but two planner paths build one over the other + * side's join key, e.g. `bucket(4, b + c)` from `b + c`. `createPartitioning` does, and so does + * a broadcast hash join, which reports the streamed side's layout over the build side's join + * key. No single cluster key stands for such an expression, so it maps to nothing. + * `canCreatePartitioning` and `areKeysCompatible` then turn the spec away. Under + * `spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled`, + * `KeyedPartitioning.createShuffleSpec` projects the position away from an unmarked layout + * where another position maps. + * + * This says which cluster key an expression is a function of, not which function it is. + * `bucket(4, b + 1)` maps to `b` just as `bucket(4, b)` does, see `argumentsAreColumns`. */ lazy val keyPositions: Seq[mutable.BitSet] = { val distKeyToPos = mutable.Map.empty[Expression, mutable.BitSet] @@ -1643,8 +1678,11 @@ case class KeyedShuffleSpec( } partitioning.expressions.map { e => val refs = e.references - assert(refs.size == 1, s"Expected exactly one child from $e, but found ${refs.size}") - distKeyToPos.getOrElse(refs.head.canonicalized, mutable.BitSet.empty) + if (refs.size == 1) { + distKeyToPos.getOrElse(refs.head.canonicalized, mutable.BitSet.empty) + } else { + mutable.BitSet.empty + } } } @@ -1659,9 +1697,10 @@ case class KeyedShuffleSpec( // 3.2 for each pair of partition expressions at the same index, the corresponding // partition keys must share overlapping positions in their respective clustering keys. // 3.3 each pair of partition expressions at the same index must describe one key space: two - // bare references, the same transform function, or the two sides of one reduce. A pair - // the join would first reduce onto one key space does not count, see - // `areKeysCompatible`'s `allowReduce`. + // bare references, the same transform function, or two transforms reduced through the + // same pairing. A transform of an expression counts only against one over the same + // argument shape. A pair the join would first reduce onto one key space does not count, + // see `areKeysCompatible`'s `allowReduce`. // 4. the partition values from both sides are following the same order. case otherSpec @ KeyedShuffleSpec(otherPartitioning, otherDistribution, _) => distribution.clustering.length == otherDistribution.clustering.length && @@ -1751,7 +1790,16 @@ case class KeyedShuffleSpec( left: Expression, right: Expression, allowReduce: Boolean): Boolean = { - if (TransformExpression.hasReducedKeys(left) || TransformExpression.hasReducedKeys(right)) { + val overExpression = Seq(left, right).exists { + case t: TransformExpression => !argumentsAreColumns(t) + case _ => false + } + if (overExpression && !argumentShape(left).exists(argumentShape(right).contains)) { + // A transform of an expression pairs only with a transform over the same argument shape. + // See `argumentsAreColumns`. + false + } else if ( + TransformExpression.hasReducedKeys(left) || TransformExpression.hasReducedKeys(right)) { // Reduced keys are in a key space that neither transform names, so comparing the transforms // says nothing about whether the two sides are laid out the same way. The pair that was // reduced together is laid out the same way, since its two sides came out of one reduce onto @@ -1776,6 +1824,48 @@ case class KeyedShuffleSpec( } } + /** + * Whether every argument of `transform` is a bare column. `isExpressionCompatible` asks it of + * both sides before it compares anything, `canCreatePartitioning` asks it of a template, and + * `rebuiltOver` asserts it. + * + * `keyPositions` pairs two sides up by the column each expression references. That pairing + * describes one key space when the argument is the column itself. `bucket(4, b + 1)` is a + * function of `b`, but not the same function of it as `bucket(4, x)` is of `x`. Nor is + * `bucket(4, s.a)` the same function of `s` as `bucket(4, t.b)` is of `t`. So such a transform + * pairs only with a transform over the same `argumentShape`. Over the same shape, the two + * compare as two transforms of columns do, e.g. `bucket(4, b + 1)` with `bucket(4, c + 1)` on a + * join `b = c`. A reduce of `bucket(8, b + 1)` onto `bucket(4, c + 1)` is fine too, since it + * maps the partition keys and evaluates no argument. An identity side is not reduced onto such a + * transform, since that would evaluate `id + 1` on keys the query never computes it for, which + * can overflow. + */ + private def argumentsAreColumns(transform: TransformExpression): Boolean = + transform.children.forall(_.isInstanceOf[Attribute]) + + /** + * `transform`, whose arguments are bare columns, with each replaced by `key`, which may be any + * expression. Any other argument would lose what surrounds its column, e.g. the `+ 1` of + * `bucket(4, b + 1)`. + */ + private def rebuiltOver(transform: TransformExpression, key: Expression) + : TransformExpression = { + assert(argumentsAreColumns(transform), s"Cannot rebuild $transform over $key") + transform.copy(children = transform.children.map(_ => key)) + } + + /** + * The arguments of a transform with their one column replaced by a placeholder of its type, or + * `None` when they reference other than one column or are not deterministic, or `expression` is + * not a transform. Two transforms with the same function and shape are the same function of + * their columns, so a join that pairs the columns up pairs the keys up too. + */ + private def argumentShape(expression: Expression): Option[Seq[Expression]] = expression match { + case t: TransformExpression if t.references.size == 1 && t.deterministic => + Some(t.children.map(QueryPlan.normalizeExpressions(_, t.references.toSeq))) + case _ => None + } + /** * Whether a join may reduce one or both sides' partition keys onto a common key space, which is * what lets two different transforms be compatible in the first place. @@ -1825,8 +1915,8 @@ case class KeyedShuffleSpec( if (otherReducer.isEmpty) { // Only this side reduces. The reducer contract is r(f1(x)) = f2(x) where "=" matches // both value and data type, so the reduced keys equal the target transform applied to - // this side's child. Report the target transform (re-targeted at the reporting - // partitioning's key attribute by `GroupPartitionsExec`) instead of the un-reduced + // this side's child. Report the target transform (rebuilt over the reporting + // partitioning's own argument by `GroupPartitionsExec`) instead of the un-reduced // `e1`, whose type can differ from the reduced keys. A connector that violates the // contract with a reducer of a different result type fails the reduced-types check in // `EnsureRequirements`, since the other side's keys are typed by the target transform. @@ -1853,15 +1943,16 @@ case class KeyedShuffleSpec( (thisResult, otherResult) // Identity transform on this side, arbitrary transform on the other side: create a reducer - // that applies the other's transform to the raw identity values. Each partition expression - // is guaranteed to have exactly one leaf child (asserted in keyPositions), which - // `IdentityReducer` binds to ordinal 0. + // that applies the other's transform to the raw identity values. `areKeysCompatible` admitted + // this pair, so `t`'s argument is one bare column. `isExpressionCompatible` refuses any other + // argument here, and `keyPositions` maps an expression with more references to no position. + // `IdentityReducer` binds `a` to ordinal 0. case (a: AttributeReference, t: TransformExpression) => - (Some(KeyReducer(IdentityReducer(t.withReference(a)), t)), None) + (Some(KeyReducer(IdentityReducer(rebuiltOver(t, a)), t)), None) // Symmetric: identity transform on the other side. case (t: TransformExpression, a: AttributeReference) => - (None, Some(KeyReducer(IdentityReducer(t.withReference(a)), t))) + (None, Some(KeyReducer(IdentityReducer(rebuiltOver(t, a)), t))) case (_, _) => (None, None) } @@ -1888,8 +1979,14 @@ case class KeyedShuffleSpec( // cannot rewrite it. This also keeps the unprojected spec returned by `createShuffleSpec` // for a marked narrowing projection from being chosen as the best spec. keyPositions.forall(_.nonEmpty) && - partitioning.expressions.forall { e => - e.isInstanceOf[AttributeReference] || e.isInstanceOf[TransformExpression] + partitioning.expressions.forall { + case _: AttributeReference => true + // `createPartitioning` replaces a transform's argument with the other child's cluster key, + // which drops whatever surrounds the key, e.g. the `+ 1` of `bucket(4, b + 1)` or the `.a` + // of `bucket(4, s.a)`. Such a spec's own child would then be laid out differently from the + // children shuffled onto it. + case t: TransformExpression => argumentsAreColumns(t) + case _ => false } && // Shuffling another child onto these keys evaluates the partition expressions per row to // decide where each row goes, and reduced keys are not what those expressions compute. @@ -1901,7 +1998,8 @@ case class KeyedShuffleSpec( val newExpressions = partitioning.expressions.zip(keyPositions).map { case (te: TransformExpression, positionSet) => - te.copy(children = te.children.map(_ => clustering(positionSet.head))) + // `canCreatePartitioning` admits only a template over bare columns, see there. + rebuiltOver(te, clustering(positionSet.head)) case (_, positionSet) => clustering(positionSet.head) } // The shuffled side is laid out on this side's partition keys, so it inherits the flag. That @@ -1968,7 +2066,7 @@ case class ShuffleSpecCollection(specs: Seq[ShuffleSpec]) extends ShuffleSpec { } override def canCreatePartitioning: Boolean = - specs.forall(_.canCreatePartitioning) + specs.exists(_.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 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 15e7f3ad41d06..1dc93d2fe1b44 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, AttributeReference, DirectShufflePartitionID, Expression, TransformExpression} +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, DirectShufflePartitionID, Expression, GetStructField, 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.{FlipLowBitFunction, Reducer, ReducibleFunction, ScalarFunction} @@ -685,6 +685,111 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper { } } + test("SPARK-59887: a spec over a transform of an expression pairs only with the same shape") { + // A join pairs its two sides up by the column each partition expression references. Over + // `a = b` it would pair `bucket(4, a + 1)` with `bucket(4, b)`, and over `s = t` it would pair + // `bucket(4, s.x)` with `bucket(4, t.y)`. A row with `a = b` sits in a different bucket on each + // side. `bucket(4, a + 1)` and `bucket(4, b + 1)` are the same function of `a` and of `b`, so + // they pair up. + val fn = new FakeBucket + val a = $"a".long + val b = $"b".long + val structType = new StructType().add("x", LongType).add("y", LongType) + val s = $"s".struct(structType) + val t = $"t".struct(structType) + // A spec clustered on the column its one partition expression references. + def keyed(expression: Expression): KeyedShuffleSpec = + KeyedShuffleSpec( + KeyedPartitioning(Seq(expression), Seq(InternalRow(0L), InternalRow(1L))), + ClusteredDistribution(expression.references.toSeq)) + def bucket(numBuckets: Int, argument: Expression): TransformExpression = + TransformExpression(fn, Seq(argument), Some(numBuckets)) + def spec(numBuckets: Int, argument: Expression): KeyedShuffleSpec = + keyed(bucket(numBuckets, argument)) + // Whether the two sides of one reduce, `bucket(8, left)` with `bucket(4, right)`, pair up. + def pairedByReduce(left: Expression, right: Expression): Boolean = + keyed(bucket(8, left).reducedTogetherWith(bucket(4, right))) + .isCompatibleWith(keyed(bucket(4, right).reducedTogetherWith(bucket(8, left)))) + val column = spec(4, b) + + withSQLConf( + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false", + SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") { + assert(spec(4, a).isCompatibleWith(column), "two columns pair up") + assert(spec(4, a).canCreatePartitioning) + assert(spec(8, a).areKeysCompatible(column, allowReduce = true), + "the fixture reduces two columns") + assert(pairedByReduce(a, b), "two columns reduced together share one key space") + assert(keyed(b).areKeysCompatible(spec(4, a), allowReduce = true), + "an identity side reduces onto a column's transform") + + Seq( + ("an expression", a + 1L, b + 1L, b + 2L), + ("a struct field", GetStructField(s, 0), GetStructField(t, 0), GetStructField(t, 1)) + ).foreach { case (shape, argument, sameShape, otherShape) => + val over = spec(4, argument) + assert(over.isCompatibleWith(spec(4, sameShape)), s"over $shape, the same shape") + assert(over.isCompatibleWith(over), s"over $shape, with itself") + assert(!over.isCompatibleWith(spec(4, otherShape)), s"over $shape, another shape") + assert(!over.isCompatibleWith(column), s"over $shape, against a column") + assert(!column.isCompatibleWith(over), s"over $shape, the other way round") + assert(!over.canCreatePartitioning, s"over $shape, no layout to shuffle onto") + + // Over the same shape, such a pair reduces as two columns do, whichever side reduces. It + // does not reduce onto another shape or a column, and an identity side does not reduce + // onto it. + assert(spec(8, argument).areKeysCompatible(spec(4, sameShape), allowReduce = true), + s"over $shape, reducing the same shape") + assert(spec(4, sameShape).areKeysCompatible(spec(8, argument), allowReduce = true), + s"over $shape, reducing the same shape the other way round") + assert(!spec(8, argument).areKeysCompatible(spec(4, otherShape), allowReduce = true), + s"over $shape, reducing another shape") + assert(!spec(8, argument).areKeysCompatible(column, allowReduce = true), s"over $shape") + assert(!column.areKeysCompatible(spec(8, argument), allowReduce = true), s"over $shape") + assert(!over.areKeysCompatible(spec(8, b), allowReduce = true), s"over $shape") + assert(!keyed(b).areKeysCompatible(over, allowReduce = true), + s"over $shape, reducing an identity side") + // Keys an earlier join reduced are in the space of that reduce. Two transforms reduced + // through the same pairing pair over the same shape, as two columns do. `bucket(12)` + // reduced with `bucket(8)` holds buckets of 4, with `bucket(18)` buckets of 6, and + // unreduced buckets of 12. + assert(pairedByReduce(argument, sameShape), s"over $shape, reduced together") + assert(!pairedByReduce(argument, otherShape), + s"over $shape, another shape reduced together") + val reducedWith8 = keyed(bucket(12, argument).reducedTogetherWith(bucket(8, a))) + assert(!reducedWith8.isCompatibleWith( + keyed(bucket(12, sameShape).reducedTogetherWith(bucket(18, b)))), + s"over $shape, reduced through different pairings") + assert(!reducedWith8.isCompatibleWith(spec(12, sameShape)), + s"over $shape, reduced against unreduced") + assert(!spec(12, sameShape).isCompatibleWith(reducedWith8), + s"over $shape, unreduced against reduced") + + // One such member does not keep a sibling from serving as the layout. + assert(ShuffleSpecCollection(Seq(over, spec(4, a))).canCreatePartitioning) + } + } + } + + test("SPARK-59901: an expression over two columns maps to no cluster key") { + // A one-side shuffle builds its partition expression over the other side's join key, which can + // reference two columns. No single cluster key stands for it. + val a = $"a".long + val b = $"b".long + Seq(a + b, TransformExpression(new FakeBucket, Seq(a + b), Some(4))).foreach { e => + val twoColumns = KeyedShuffleSpec( + KeyedPartitioning(Seq(e), Seq(InternalRow(0L), InternalRow(1L))), + ClusteredDistribution(Seq(a, b))) + assert(twoColumns.keyPositions.forall(_.isEmpty), s"$e") + assert(!twoColumns.isCompatibleWith(twoColumns), s"$e") + withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") { + assert(!twoColumns.canCreatePartitioning, s"$e") + } + } + } + test("createShuffleSpec: a marked narrowing projection yields an unusable spec") { val a = $"a".int val b = $"b".int 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 421a635bc33b6..ab69926254cf3 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 @@ -76,8 +76,8 @@ case class GroupPartitionsExec( // There can be multiple `KeyedPartitioning`s in an output partitioning of a join, but they // can only differ in `expressions`. Their `partitionKeys` reference and `isCollapsed` flag // are shared (enforced by `PartitioningCollection`), so the grouping is computed once. - // When reducers are applied, the stored reduced expressions are re-targeted at each - // `KeyedPartitioning`'s own key attribute and reported instead of the original ones. Their + // When reducers are applied, the stored reduced expressions are rebuilt over each + // `KeyedPartitioning`'s own argument and reported instead of the original ones. Their // data types match the reduced partition keys for the identity-vs-transform and // single-side-transform reducers; for the both-sides-reduce shape no single transform // describes the keys (see `KeyedShuffleSpec.reducersBothWays`). @@ -107,9 +107,16 @@ case class GroupPartitionsExec( 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.withReference(expr.references.head) + // picked (`collectFirst`). The keys are reduced once, from the shared key rows, + // so `reduced` describes them whichever member this is. Only its function says + // that, so the member's own argument takes the place of `reduced`'s. A side + // shuffled onto this layout reports `bucket(8, b + 1)` next to `bucket(8, id)`, + // or `b + 1` next to an identity `id`. + val argument = expr match { + case t: TransformExpression => t.children + case e => Seq(e) + } + reduced.copy(children = argument) case (expr, None) => expr } case None => projectedExpressions 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 7d2d1228b63b9..32b5b58cd816b 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 @@ -206,7 +206,11 @@ case class EnsureRequirements( val shouldConsiderMinParallelism = specs.forall(p => !p._2.canCreatePartitioning || children(p._1).isInstanceOf[ShuffleExchangeLike] ) - // Choose all the specs that can be used to shuffle other children + // Choose all the specs that can be used to shuffle other children. For a collection only the + // members that can serve as the layout count, here and below. Any other member would build a + // partitioning its own child is not laid out on. + def servingMembers(spec: ShuffleSpec): Seq[ShuffleSpec] = + flattenSpec(spec).filter(_.canCreatePartitioning) val candidateSpecs = specs.filter { case (index, spec) => spec.canCreatePartitioning && (!shouldConsiderMinParallelism || @@ -228,8 +232,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. A collection counts with its first member that + // can serve as the layout, which is its first member unless that one cannot. + Some(finalCandidateSpecs.values.maxBy(servingMembers(_).head.numPartitions)) } // Check if the following conditions are satisfied: @@ -263,14 +268,14 @@ case class EnsureRequirements( // 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))) + childrenIndexes.filter(i => servingMembers(best).exists(_.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) + servingMembers(best) .filter(m => matchedMembers.forall(_.exists(m.isCompatibleWith))) .maxByOption(_.numPartitions) } @@ -929,15 +934,17 @@ case class EnsureRequirements( partitioning: Partitioning, distribution: ClusteredDistribution): Option[KeyedShuffleSpec] = { def tryCreate(partitioning: KeyedPartitioning): Option[KeyedShuffleSpec] = { - // The single-column invariant in KeyedPartitioning.supportsExpressions guarantees one - // attribute per partition expression. val attributes = partitioning.expressions.flatMap(_.references) val clustering = distribution.clustering val satisfies = if (SQLConf.get.getConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION)) { - attributes.length == clustering.length && attributes.zip(clustering).forall { - case (l, r) => l.semanticEquals(r) - } + // Only an expression over a single column covers that column. One over several, e.g. + // `b + c`, maps to no position (`KeyedShuffleSpec.keyPositions`). The spec turns it away + // or projects it away, so a layout with one does not cover the cluster keys. + partitioning.expressions.forall(_.references.size == 1) && + attributes.length == clustering.length && attributes.zip(clustering).forall { + case (l, r) => l.semanticEquals(r) + } } else { partitioning.satisfies(distribution) } 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 38372fce57753..63f5df018c72a 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 @@ -44,7 +44,7 @@ import org.apache.spark.sql.execution.{ UnionExec} import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation, GroupPartitionsExec} import org.apache.spark.sql.execution.exchange.{ShuffleExchangeExec, ShuffleExchangeLike, ValidateRequirements} -import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, ShuffledHashJoinExec, SortMergeJoinExec} import org.apache.spark.sql.functions.{col, max} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.SQLConf._ @@ -327,6 +327,33 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with private def createBucketedIdTables(bucketCounts: Int*): Unit = bucketCounts.foreach(n => createBucketedIdTable(s"bucket$n", n)) + /** Creates `ident`, identity-partitioned on `id`, holding `ids`. */ + private def createIdentityIdTable(ids: Seq[Long] = 0L until 8L): Unit = { + createTable("ident", Array(Column.create("id", LongType)), Array(identity("id"))) + sql("INSERT INTO testcat.ns.ident VALUES " + ids.map(i => s"($i)").mkString(", ")) + } + + /** Creates `name`, unpartitioned, with a column `b` holding `bs`. */ + private def createPlainTable(bs: Seq[Long] = 0L until 8L, name: String = "plain"): Unit = { + createTable(name, Array(Column.create("b", LongType)), Array.empty) + sql(s"INSERT INTO testcat.ns.$name VALUES " + bs.map(b => s"($b)").mkString(", ")) + } + + /** Creates `plain2`, unpartitioned, with columns `b` and `c`, where `b + c` runs from 1 to 3. */ + private def createPlain2Table(): Unit = { + createTable("plain2", Array(Column.create("b", LongType), Column.create("c", LongType)), + Array.empty) + sql("INSERT INTO testcat.ns.plain2 VALUES (0, 1), (0, 2), (1, 1), (1, 2)") + } + + /** Creates `ps` with a struct column `s` of 8 rows, `(a = i, b = 7 - i)`. */ + private def createStructTable(): Unit = { + createTable("ps", Array(Column.create("s", new StructType().add("a", LongType) + .add("b", LongType))), Array.empty) + sql("INSERT INTO testcat.ns.ps VALUES " + + (0 until 8).map(i => s"(named_struct('a', ${i}L, 'b', ${7 - i}L))").mkString(", ")) + } + /** Joins `bucket12`, `bucket8` and `bucket` on `id`, in that order. */ private def threeWayBucketJoinDF(third: Int): DataFrame = sql("SELECT b12.id FROM testcat.ns.bucket12 b12 " + @@ -946,6 +973,409 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with } } + // Each key is a function of `b` other than `b` itself. A check on the leaves of the key would + // reject `b + 1` for its literal, but not `-b`, so both are tested. + Seq("p.b + 1" -> ((b: Long) => b + 1), "-p.b" -> ((b: Long) => -b)).foreach { + case (key, keyOf) => + // `plain` holds `b` in 0..7, and `keys`, bucketed on `id`, holds the key of each `b`. + def createKeyTables(): Unit = { + createPlainTable() + createTable("keys", Array(Column.create("id", LongType)), Array(bucket(4, "id"))) + sql("INSERT INTO testcat.ns.keys VALUES " + + (0 until 8).map(b => s"(${keyOf(b)})").mkString(", ")) + } + + test("SPARK-59887: a side shuffled onto a transform of an expression is not paired as it " + + s"is ($key)") { + createKeyTables() + Seq(4, 8).foreach(n => createBucketedIdTable(s"bucket$n", n, numIds = 8)) + + // The first join shuffles `plain` onto the bucket of its key. That is a function of `b`, + // but not the same function of it as the third table's `bucket(n, id)` is of `id`. The + // second join clusters `plain` on `b`. Pairing the two there, as they stand or after + // reducing `bucket(8, id)` onto 4 buckets, puts a row with `b = id` on a different + // partition on each side. + Seq(4 -> "false", 8 -> "true").foreach { case (third, allowCompatibleTransforms) => + withSQLConf( + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> allowCompatibleTransforms) { + val df = sql( + s"""SELECT k.id, p.b, t.id FROM testcat.ns.keys k + |JOIN testcat.ns.plain p ON k.id = $key + |JOIN testcat.ns.bucket$third t ON p.b = t.id""".stripMargin) + checkAnswer(df, (0 until 8).map(b => Row(keyOf(b), b.toLong, b.toLong))) + // One shuffle per join: `plain` for the first, the first join's output for the + // second, and none for the third table. The second join does not shuffle onto the + // bucket of the key, since that would move both of its sides. + assert(collectShuffles(stripAQEPlan(df.queryExecution.executedPlan)).size == 2) + } + } + } + + test(s"SPARK-59887: a side is not shuffled onto a transform of an expression ($key)") { + createKeyTables() + createTable("xy", Array(Column.create("x", LongType), Column.create("y", LongType)), + Array.empty) + sql("INSERT INTO testcat.ns.xy VALUES " + + (0 until 8).map(b => s"(${keyOf(b)}, $b)").mkString(", ")) + + // The first join shuffles `plain` onto the bucket of its key, so its output offers that + // layout next to `bucket(4, id)`. The second join must not shuffle `xy` onto the first + // one. That builds `bucket(4, y)` for `xy`, not the bucket of the same function of `y`. A + // row with `b = y` then lands in a different bucket on each side. The join order decides + // which of the two layouts comes first, so both orders run. + withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") { + Seq( + s"testcat.ns.plain p JOIN testcat.ns.keys k ON $key = k.id", + s"testcat.ns.keys k JOIN testcat.ns.plain p ON k.id = $key").foreach { firstJoin => + val df = sql( + s"""SELECT k.id, p.b, q.x FROM $firstJoin + |JOIN testcat.ns.xy q ON k.id = q.x AND p.b = q.y""".stripMargin) + checkAnswer(df, (0 until 8).map(b => Row(keyOf(b), b.toLong, keyOf(b)))) + // `xy` is shuffled onto the usable layout, `bucket(4, x)`, and never onto a bucket of + // `y`. So `plain` and `xy` are the only sides that move. + val shuffles = collectShuffles(stripAQEPlan(df.queryExecution.executedPlan)) + assert(shuffles.size == 2) + assert(!keyedPartitioningsOf(shuffles).exists(_.references.exists(_.name == "y"))) + } + } + } + + test("SPARK-59887: a broadcast join's layout over a join key is not paired as it is " + + s"($key)") { + createKeyTables() + createBucketedIdTable("bucket4", 4, numIds = 8) + + // An inner broadcast join reports the streamed side's layout over the build side's join + // key too, so the first join also offers the bucket of `plain`'s key. That needs no conf + // beyond the defaults. The second join clusters `plain` on `b`, so pairing that layout with + // `bucket(4, id)` as it stands puts a row with `b = id` on a different partition on each + // side. + val df = sql( + s"""SELECT /*+ BROADCAST(p) */ k.id, p.b, t.id FROM testcat.ns.keys k + |JOIN testcat.ns.plain p ON k.id = $key + |JOIN testcat.ns.bucket4 t ON p.b = t.id""".stripMargin) + checkAnswer(df, (0 until 8).map(b => Row(keyOf(b), b.toLong, b.toLong))) + assert(collect(df.queryExecution.executedPlan) { case j: BroadcastHashJoinExec => j } + .nonEmpty, "the first join is a broadcast join") + } + } + + test("SPARK-59887: a layout over a transform of a struct field is not paired nor shuffled onto") { + createBucketedIdTable("bucket4", 4, numIds = 8) + createStructTable() + // `l` shuffles `ps` onto `bucket(4, s.a)`, and `r` shuffles it onto `bucket(4, s.b)`. + val l = "SELECT x.id, ps.s FROM testcat.ns.bucket4 x JOIN testcat.ns.ps ps ON x.id = ps.s.a" + val r = "SELECT y.id, qs.s FROM testcat.ns.bucket4 y JOIN testcat.ns.ps qs ON y.id = qs.s.b" + + withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") { + // Both layouts are functions of `s`, but not the same one, so the join on `s` must not pair + // them as they stand. A row lands in bucket `a % 4` on one side and `b % 4` on the other. + checkAnswer( + sql(s"SELECT l.id, r.id FROM ($l) l JOIN ($r) r ON l.s = r.s"), + (0 until 8).map(i => Row(i.toLong, (7 - i).toLong))) + // Nor may a join shuffle a plain `ps` onto `l`'s layout. That builds `bucket(4, s)` for it, + // which drops the `.a` and hands the whole struct to a bucket function over a long. + checkAnswer( + sql(s"SELECT l.id FROM ($l) l JOIN testcat.ns.ps r ON l.s = r.s"), + (0 until 8).map(i => Row(i.toLong))) + } + } + + test("SPARK-59887: a reduce keeps what each member's keys are computed from") { + Seq(4, 2).foreach(n => createBucketedIdTable(s"bucket$n", n, numIds = 8)) + createIdentityIdTable() + createPlainTable() + createStructTable() + + withSQLConf( + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") { + // The first join shuffles `plain` onto `bucket(4, b + 1)`, or onto `b + 1` for the + // identity-partitioned `ident`. The second join reduces the first one's output onto + // `bucket(2)`, `plain`'s member included. That member is still over `b + 1`. So the third + // join, on the bare `b`, must neither pair it with `bucket(2, id)` nor shuffle a plain + // side onto it. + for { + first <- Seq("bucket4", "ident") + (third, thirdKey) <- Seq("bucket2" -> "id", "plain" -> "b") + } { + val df = sql( + s"""SELECT f.id, p.b, t2.id, t3.$thirdKey FROM testcat.ns.$first f + |JOIN testcat.ns.plain p ON f.id = p.b + 1 + |JOIN testcat.ns.bucket2 t2 ON f.id = t2.id + |JOIN testcat.ns.$third t3 ON p.b = t3.$thirdKey""".stripMargin) + checkAnswer(df, (0 until 7).map(b => Row(b + 1L, b.toLong, b + 1L, b.toLong))) + } + + // The same for a struct field. `ps`'s member is still over `s.a` after the reduce, so the + // third join must not shuffle a plain `ps` onto it, which would build `bucket(2, s)`. + val df = sql( + """SELECT x.id FROM testcat.ns.bucket4 x + |JOIN testcat.ns.ps ps ON x.id = ps.s.a + |JOIN testcat.ns.bucket2 t2 ON x.id = t2.id + |JOIN testcat.ns.ps r ON ps.s = r.s""".stripMargin) + checkAnswer(df, (0 until 8).map(i => Row(i.toLong))) + } + } + + // `ident` also holds an id that the key's function overflows on under ANSI. + Seq( + ("p.b + 1", 0L until 8L, Long.MaxValue, (b: Long) => b + 1), + ("-p.b", -7L to 0L, Long.MinValue, (b: Long) => -b) + ).foreach { case (key, bs, overflowingId, keyOf) => + test(s"SPARK-59887: an identity side is not reduced onto a transform of an expression ($key)") { + createBucketedIdTable("bucket4", 4, numIds = 8) + createPlainTable(bs) + createIdentityIdTable(bs :+ overflowingId) + + // The first join shuffles `plain` onto the bucket of its key. Reducing `ident` onto it would + // evaluate the key's function on `ident`'s partition keys while planning, which overflows for + // that id. The query itself never computes it, since no `b` matches that id. + withSQLConf( + SQLConf.ANSI_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") { + val df = sql( + s"""SELECT k.id, p.b, i.id FROM testcat.ns.bucket4 k + |JOIN testcat.ns.plain p ON k.id = $key + |JOIN testcat.ns.ident i ON p.b = i.id""".stripMargin) + checkAnswer(df, + bs.filter(b => (0L until 8L).contains(keyOf(b))).map(b => Row(keyOf(b), b, b))) + } + } + } + + test("SPARK-59887: a reduce keeps a member over two columns") { + createIdentityIdTable() + createBucketedIdTable("bucket2", 2, numIds = 8) + createPlain2Table() + + // The broadcast join reports `ident`'s layout over `p.b + p.c` too, and the second join reduces + // it onto `bucket(2)`. That member has to stay over `p.b + p.c`. Reported over `p.b`, it would + // serve the `GROUP BY p.b` as it stands, while the rows of one `b` sit in two buckets. + withSQLConf(SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") { + checkAnswer( + sql("""SELECT /*+ BROADCAST(p) */ p.b, max(p.c) FROM testcat.ns.ident i + |JOIN testcat.ns.plain2 p ON i.id = p.b + p.c + |JOIN testcat.ns.bucket2 t2 ON i.id = t2.id + |GROUP BY p.b""".stripMargin), + Seq(Row(0L, 2L), Row(1L, 2L))) + } + } + + test("SPARK-59887: two layouts over the same shape of a join key pair as two columns do") { + Seq("orders", "returns").foreach(createBucketedIdTable(_, 16, numIds = 16)) + createTable("customers", Array(Column.create("id", IntegerType)), Array.empty) + sql("INSERT INTO testcat.ns.customers VALUES " + (0 until 16).map(i => s"($i)").mkString(", ")) + + // Each broadcast join reports its fact table's layout over the dimension's key too, widened to + // the fact key's type, `bucket(16, cast(c.id as bigint))`. That needs no conf beyond the + // defaults. The join on `c1.id = c2.id` pairs the two as they stand, since they are the same + // function of their columns. + val dimensions = sql( + """SELECT a.id FROM + | (SELECT /*+ BROADCAST(c1) */ c1.id FROM testcat.ns.orders o + | JOIN testcat.ns.customers c1 ON o.id = c1.id) a + |JOIN + | (SELECT /*+ BROADCAST(c2) */ c2.id FROM testcat.ns.returns r + | JOIN testcat.ns.customers c2 ON r.id = c2.id) b + |ON a.id = b.id""".stripMargin) + checkAnswer(dimensions, (0 until 16).map(Row(_))) + assert(collectAllShuffles(dimensions.queryExecution.executedPlan).isEmpty) + + // The same for two one-side shuffles onto `bucket(n, p.b + 1)` and `bucket(4, q.b + 1)`, + // joined on `p.b = q.b`. With `bucket(4)` on both sides they pair as they stand. With + // `bucket(8)` on one side, `allowCompatibleTransforms` reduces it onto `bucket(4)`. Either way + // only the two one-side shuffles are left. + Seq(4, 8).foreach(n => createBucketedIdTable(s"bucket$n", n, numIds = 8)) + createPlainTable() + // A second unpartitioned table, so that the two one-side shuffles are not one reused exchange. + createPlainTable(name = "plainq") + Seq(4, 8).foreach { n => + withSQLConf( + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") { + val legs = sql( + s"""SELECT l.b FROM + | (SELECT p.b FROM testcat.ns.bucket$n t1 + | JOIN testcat.ns.plain p ON t1.id = p.b + 1) l + |JOIN + | (SELECT q.b FROM testcat.ns.bucket4 t2 + | JOIN testcat.ns.plainq q ON t2.id = q.b + 1) r + |ON l.b = r.b""".stripMargin) + checkAnswer(legs, (0 until 7).map(b => Row(b.toLong))) + assert(collectAllShuffles(legs.queryExecution.executedPlan).size == 2, s"bucket$n") + } + } + } + + test("SPARK-59887: the same shape reduced together pairs only through the same pairing") { + Seq(12, 8, 18).foreach(n => createBucketedIdTable(s"bucket$n", n, numIds = 24)) + createPlainTable(0L until 24L) + // A second unpartitioned table, so that the two one-side shuffles are not one reused exchange. + createPlainTable(0L until 24L, name = "plainq") + + // Each leg shuffles its plain table onto `bucket(12, b + 1)`. The first leg then reduces + // `bucket(12)` with `bucket(8)`, so its `b + 1` member holds buckets of 4. The second leg + // reduces with `bucket(8)` too, with `bucket(18)`, or not at all. Only with `bucket(8)` does + // it hold buckets of 4 as well, and only then does the final join pair the two as they stand. + // Otherwise it holds buckets of 6 or of 12, and the final join shuffles both sides. + withSQLConf( + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") { + Seq(Some(8) -> 2, Some(18) -> 4, None -> 4).foreach { case (other, shuffles) => + val reduce = other.fold("")(n => s"JOIN testcat.ns.bucket$n u ON t.id = u.id") + val df = sql( + s"""SELECT l.b FROM + | (SELECT p.b FROM testcat.ns.bucket12 t JOIN testcat.ns.plain p ON t.id = p.b + 1 + | JOIN testcat.ns.bucket8 u ON t.id = u.id) l + |JOIN + | (SELECT q.b FROM testcat.ns.bucket12 t JOIN testcat.ns.plainq q ON t.id = q.b + 1 + | $reduce) r + |ON l.b = r.b""".stripMargin) + checkAnswer(df, (0 until 23).map(b => Row(b.toLong))) + assert(collectAllShuffles(df.queryExecution.executedPlan).size == shuffles, s"$other") + } + } + } + + test("SPARK-59887: a member that cannot serve does not rank its collection") { + createTable("t1", Array(Column.create("a", LongType), Column.create("c", LongType)), + Array(identity("a"), identity("c"))) + sql("INSERT INTO testcat.ns.t1 VALUES (1, 1), (1, 2), (1, 3), (2, 4), (2, 5)") + createTable("p", Array(Column.create("y", LongType)), Array.empty) + sql("INSERT INTO testcat.ns.p VALUES (0), (1), (2), (3), (4)") + createTable("q", Array(Column.create("m", LongType), Column.create("n", LongType)), + Array(identity("m"))) + sql("INSERT INTO testcat.ns.q VALUES (1, 0), (2, 3), (3, 9)") + + // The broadcast join reports `[a, c]` and `[a, y + 1]`. The first projects to `[a]`, 2 + // partitions, and can serve as the layout. The second keeps 5 partitions and cannot. So `q`, + // with 3 partitions and no exchange, ranks first, and the other side is shuffled onto it. + withSQLConf( + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val df = sql( + """SELECT /*+ BROADCAST(p) */ * FROM testcat.ns.t1 + |JOIN testcat.ns.p ON t1.c = p.y + 1 + |JOIN testcat.ns.q ON t1.a = q.m AND p.y = q.n""".stripMargin) + checkAnswer(df, Seq(Row(1L, 1L, 0L, 1L, 0L), Row(2L, 4L, 3L, 2L, 3L))) + assert(collectAllShuffles(df.queryExecution.executedPlan) + .map(_.outputPartitioning.numPartitions) == Seq(3)) + } + } + + test("SPARK-59901: a join key over two columns is not paired through one of them") { + createIdentityIdTable() + createBucketedIdTable("bucket4", 4, numIds = 8) + createPlain2Table() + val expected = Seq(Row(1L, 0L, 1L), Row(2L, 0L, 2L), Row(2L, 1L, 1L), Row(3L, 1L, 2L)) + + // A one-side shuffle builds its partition expression over the other side's join key. That is + // `b + c` here, or `bucket(4, b + c)` over a bucketed side. No single cluster key stands for + // either, so a spec over one pairs with nothing and is no layout to shuffle onto. On master, + // AQE validates the first plan and fails there, but not on this branch. The second join of + // the second query asks what to shuffle onto. + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") { + checkAnswer( + sql("""SELECT i.id, p.b, p.c FROM testcat.ns.ident i + |JOIN testcat.ns.plain2 p ON i.id = p.b + p.c""".stripMargin), + expected) + checkAnswer( + sql("""SELECT b4.id, q.b, q.c FROM testcat.ns.bucket4 b4 + |JOIN testcat.ns.plain2 p ON b4.id = p.b + p.c + |JOIN testcat.ns.plain2 q ON p.b = q.b AND p.c = q.c""".stripMargin), + expected) + } + } + + test("SPARK-59901: a join key over two columns does not cover its columns") { + Seq("ident2a", "ident2b").foreach { name => + createTable(name, Array(Column.create("x", LongType), Column.create("d", LongType)), + Array(identity("x"), identity("d"))) + sql(s"INSERT INTO testcat.ns.$name VALUES (1, 1), (2, 1), (3, 2), (4, 2)") + } + Seq("pa", "pb").foreach { name => + createTable(name, Array(Column.create("b", LongType), Column.create("c", LongType), + Column.create("d", LongType)), Array.empty) + sql(s"INSERT INTO testcat.ns.$name VALUES (0, 1, 1), (1, 1, 1), (1, 2, 2), (2, 2, 2)") + } + + // Each broadcast join reports its identity side's layout over `[b + c, d]` too. A projection + // to the join keys drops `b + c`, which maps to no position, so it does not cover `b` and `c` + // for `spark.sql.requireAllClusterKeysForCoPartition`. Otherwise the join on `(b, c, d)` would + // be paired on `d` alone. + withSQLConf( + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + val df = sql( + """SELECT l.b, l.c, l.d FROM + | (SELECT /*+ BROADCAST(p) */ p.b, p.c, p.d FROM testcat.ns.ident2a i + | JOIN testcat.ns.pa p ON i.x = p.b + p.c AND i.d = p.d) l + |JOIN + | (SELECT /*+ BROADCAST(q) */ q.b, q.c, q.d FROM testcat.ns.ident2b j + | JOIN testcat.ns.pb q ON j.x = q.b + q.c AND j.d = q.d) r + |ON l.b = r.b AND l.c = r.c AND l.d = r.d""".stripMargin) + checkAnswer(df, Seq(Row(0L, 1L, 1L), Row(1L, 1L, 1L), Row(1L, 2L, 2L), Row(2L, 2L, 2L))) + assert(collectAllGroupPartitions(df.queryExecution.executedPlan).isEmpty) + + // With the identity side's columns kept, a sibling layout such as `[x, d]` satisfies the + // join's distribution, so no leg is shuffled first. Then only the coverage check keeps + // `[b + c, d]` from pairing on `d` alone. + val kept = sql( + """SELECT l.b, l.c, l.d, l.x, l.id, r.x, r.id FROM + | (SELECT /*+ BROADCAST(p) */ p.b, p.c, p.d, i.x, i.d AS id FROM testcat.ns.ident2a i + | JOIN testcat.ns.pa p ON i.x = p.b + p.c AND i.d = p.d) l + |JOIN + | (SELECT /*+ BROADCAST(q) */ q.b, q.c, q.d, j.x, j.d AS id FROM testcat.ns.ident2b j + | JOIN testcat.ns.pb q ON j.x = q.b + q.c AND j.d = q.d) r + |ON l.b = r.b AND l.c = r.c AND l.d = r.d""".stripMargin) + checkAnswer(kept, Seq(Row(0L, 1L, 1L, 1L, 1L, 1L, 1L), Row(1L, 1L, 1L, 2L, 1L, 2L, 1L), + Row(1L, 2L, 2L, 3L, 2L, 3L, 2L), Row(2L, 2L, 2L, 4L, 2L, 4L, 2L))) + assert(collectAllGroupPartitions(kept.queryExecution.executedPlan).isEmpty) + } + } + + test("SPARK-59901: a join keyed on a join-key expression does not run in one partition") { + createTable("i1", Array(Column.create("id", LongType)), Array(identity("id"))) + sql("INSERT INTO testcat.ns.i1 VALUES (1), (2), (3), (4)") + createTable("pa", Array(Column.create("b", LongType), Column.create("c", LongType)), + Array.empty) + sql("INSERT INTO testcat.ns.pa VALUES (0, 1), (1, 1), (1, 2), (2, 2)") + createTable("rk", Array(Column.create("k", LongType)), Array.empty) + sql("INSERT INTO testcat.ns.rk VALUES (1), (2), (3), (4)") + + // The broadcast join reports `i1`'s layout over its join key, e.g. `[b + c]`. The second join + // is keyed on that expression itself, so under `requireAllClusterKeysForDistribution` the + // layout satisfies it, but no position maps to a cluster key. Projecting onto the positions + // that map would run the join in one partition. `b + 1` is the same with one column. + 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_DISTRIBUTION.key -> "true", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + Seq( + ("p.b + 1", "l.b + 1", Seq(Row(0L, 1L, 1L), Row(1L, 1L, 2L), Row(1L, 2L, 2L), + Row(2L, 2L, 3L))), + ("p.b + p.c", "l.b + l.c", Seq(Row(0L, 1L, 1L), Row(1L, 1L, 2L), Row(1L, 2L, 3L), + Row(2L, 2L, 4L)))).foreach { case (innerKey, outerKey, expected) => + val df = sql( + s"""SELECT l.b, l.c, r.k FROM + | (SELECT /*+ BROADCAST(p) */ p.b, p.c FROM testcat.ns.i1 i + | JOIN testcat.ns.pa p ON i.id = $innerKey) l + |JOIN testcat.ns.rk r ON $outerKey = r.k""".stripMargin) + checkAnswer(df, expected) + val plan = df.queryExecution.executedPlan + assert(collectAllGroupPartitions(plan).isEmpty, innerKey) + assert(collectShuffles(plan).size == 2, innerKey) + } + } + } + test("SPARK-59121: two reduced partitionings are not compatible by their transforms") { // 24 ids, so that each leg's two sides really report different key sets and the join reduces // them. With 12 ids `bucket(18, id)` is the identity and the leg would be co-partitioned as it @@ -4581,21 +5011,19 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with case class Query(left: String, right: String, on: String, order: String, select: String, expected: Seq[Row], sortingNodes: Int) val ids = (94 to 100).map(i => Row(i.toLong)) - val singleColumnKeys = Seq( + val queries = Seq( Query("ident", "plain", "i.id = 100 - p.b", "100 - p.b", "p.b", (6 to 0 by -1).map(b => Row(b.toLong)), 0), Query("ident", "plain", "i.id = 100 - p.b", "100 - p.b DESC", "p.b", (0 to 6).map(b => Row(b.toLong)), 1), Query("ident", "plain", "i.id = p.s.a", "p.s.a", "i.id", ids, 0), Query("ident2", "plain2", "i.x = 100 - p.b AND i.y = p.c", "100 - p.b, p.c DESC", "p.b, p.c", - Seq(Row(6L, 2L), Row(6L, 1L), Row(5L, 1L)), 1)) - // AQE's plan validation fails on `b + c` until SPARK-59901 is fixed, so it runs with AQE off. - val twoColumnKey = - Query("ident", "plain", "i.id = p.b + p.c", "p.b + p.c", "i.id", ids, 0) + Seq(Row(6L, 2L), Row(6L, 1L), Row(5L, 1L)), 1), + Query("ident", "plain", "i.id = p.b + p.c", "p.b + p.c", "i.id", ids, 0)) for { - (query, aqe) <- - singleColumnKeys.flatMap(q => Seq(q -> true, q -> false)) :+ (twoColumnKey -> false) + query <- queries + aqe <- Seq(true, false) broadcast <- Seq(false, true) } { withSQLConf( 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 56804b74d41d3..2536e120e3932 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 @@ -1429,6 +1429,70 @@ class EnsureRequirementsSuite extends SharedSparkSession { } } + test("SPARK-59887: a collection ranks by its first member that can serve as the layout") { + val id = AttributeReference("id", IntegerType)() + val c = AttributeReference("c", IntegerType)() + val y = AttributeReference("y", IntegerType)() + val m = AttributeReference("m", IntegerType)() + val n = AttributeReference("n", IntegerType)() + // Clustered on (id, y), the first member `[id, y + 1]` keeps its five partitions but cannot + // serve as the layout, since `y + 1` is not a column. The second member `[id, c]` projects to + // `[id]`, two partitions, and can. The other side keeps three partitions and can serve too. So + // the other side wins, as it does when the collection cannot serve at all. Ranking the + // collection by its first member would let it win with five, and shuffle the other side onto + // the two partitions of the member that can serve. + val keys = Seq(InternalRow(1, 1), InternalRow(1, 2), InternalRow(1, 3), InternalRow(2, 4), + InternalRow(2, 5)) + val collection = new DummySparkPlanWithBatchScanChild( + outputPartitioning = PartitioningCollection.fromPartitionings(Seq( + KeyedPartitioning(Seq(id, Add(y, Literal(1))), keys), + KeyedPartitioning(Seq(id, c), keys)))) + val other = new DummySparkPlanWithBatchScanChild( + outputPartitioning = KeyedPartitioning(Seq(m), Seq(InternalRow(1), InternalRow(2), + InternalRow(3)))) + + withSQLConf( + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", + REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false") { + val smj = SortMergeJoinExec(Seq(id, y), Seq(m, n), Inner, None, collection, other) + val planned = EnsureRequirements.apply(smj).asInstanceOf[SortMergeJoinExec] + assert(planned.right.collect { case s: ShuffleExchangeExec => s }.isEmpty, + "the other side is the layout") + assert(planned.left.collect { case s: ShuffleExchangeExec => s } + .map(_.outputPartitioning.numPartitions) === Seq(3)) + } + } + + test("SPARK-59887: a child matched only through a member that cannot serve is shuffled") { + val id = AttributeReference("id", IntegerType)() + val b = AttributeReference("b", IntegerType)() + val z = AttributeReference("z", IntegerType)() + val c = AttributeReference("c", IntegerType)() + // Clustered on (id, b), the first child reports `bucket(4, id)`, which can serve as the + // layout, and `bucket(4, b + 1)`, which cannot. The second child reports `bucket(4, c + 1)` + // over (z, c), which pairs with the second member only. Matching it through that member would + // leave no member that serves every matched child, and both sides would be shuffled. It is + // shuffled onto `bucket(4, id)` instead, and the first child stays as it is. Without this + // change both children stay as they are, paired through `bucket(4, b + 1)`. That plan is + // correct, and giving it up is a cost the description lists. + val keys = (0 until 4).map(InternalRow(_)) + val collection = new DummySparkPlanWithBatchScanChild( + outputPartitioning = PartitioningCollection.fromPartitionings(Seq( + KeyedPartitioning(Seq(bucket(4, id)), keys), + KeyedPartitioning(Seq(bucket(4, Add(b, Literal(1)))), keys)))) + val other = new DummySparkPlanWithBatchScanChild( + outputPartitioning = KeyedPartitioning(Seq(bucket(4, Add(c, Literal(1)))), keys)) + + withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") { + val smj = SortMergeJoinExec(Seq(id, b), Seq(z, c), Inner, None, collection, other) + val planned = EnsureRequirements.apply(smj).asInstanceOf[SortMergeJoinExec] + assert(planned.left.collect { case s: ShuffleExchangeExec => s }.isEmpty, + "the first child is the layout") + assert(planned.right.collect { case s: ShuffleExchangeExec => s }.size === 1) + } + } + private class DummySparkPlanWithBatchScanChild(outputPartitioning: Partitioning) extends DummySparkPlan( children = Seq(BatchScanExec(Seq.empty, null, Seq.empty, table = null)),