From 398802fef0b3aba434235094d7d31f01b6753c42 Mon Sep 17 00:00:00 2001 From: Peter Toth Date: Sat, 3 Oct 2026 10:01:07 +0200 Subject: [PATCH] [SPARK-59887][SPARK-59901][SQL][4.3] Fix wrong results when a storage-partitioned join pairs a transform of a join-key expression with one of a column ### Differences from #59165 This is the branch-4.3 backport of #59165. The description below is the original one. On this branch: - The cherry-pick conflicts in `partitioning.scala`, `GroupPartitionsExec`, `EnsureRequirements` and the suite's imports, since the code around the change differs on branch-4.3. The main code is ported by hand. The tests are the same as on master. - Branch-4.3 has no `pickCoPartitionTarget`. `EnsureRequirements` picks the layout through `bestSpecOpt` and `bestMemberOpt` instead. Only a member that can serve as the layout matches a child there, and `bestMemberOpt` picks only such a member. - Branch-4.3 ranks a collection by its first member. Now that a collection can serve when any member can, that first member may be one that cannot. So a collection ranks by its first member that can serve, which is the branch's ranking whenever the first member can. - Two new `EnsureRequirementsSuite` tests cover these two changes, since no master test reaches them on this branch. "SPARK-59887: a collection ranks by its first member that can serve as the layout" lists a member that cannot serve first. "SPARK-59887: a child matched only through a member that cannot serve is shuffled" pairs a child with the layout's other member only. Each fails with its change removed. - The `spark.sql.requireAllClusterKeysForCoPartition` change is left out. Here a member over two columns never satisfies a clustered distribution under `allowKeysSubsetOfPartitionKeys`, so its child is shuffled before the co-partitioning check. Its test, "SPARK-59901: a join key over two columns does not cover its columns", is kept, and it passes on this branch without this change. - `TransformExpressionSuite` does not exist on branch-4.3, so its hunk is left out. So are the doc hunks for code that branch-4.3 does not have. - The bugs reproduce on branch-4.3 as on master, `b + 1` included, through both the one-side shuffle and the broadcast join. Measured without this change: - the `b + 1` query above returns 0 of 7 rows; - the `-p.b` shapes return 4 of 8 rows; - the struct field pairing returns 0 of 8 rows, and shuffling onto it fails with the `ClassCastException`. - New tests: all but three fail on branch-4.3 without this change. The three are the two that also pass on master, and the coverage test above. Both new `ShuffleSpecSuite` tests fail too. The SPARK-59905 test's `p.b + p.c` run with AQE on passes without this change as well. - The main changes were removed one at a time: the shape gate, the `canCreatePartitioning` clause, the `forall`, the member filter, the ranking change, the matching change, the `GroupPartitionsExec` change and the `keyPositions` assertion. Each fails the same tests as on master, except: - without the member filter, only the two `xy` tests fail, on the assertion in `rebuiltOver`; - with the `keyPositions` assertion back, only "SPARK-59901: a join key over two columns is not paired through one of them" fails among the `KeyGroupedPartitioningSuite` tests; - removing the ranking change or the matching change fails only its new `EnsureRequirementsSuite` test. The master ranking test's collection lists its serving member first, so the branch's own ranking already gets that one right. - Ran `ShuffleSpecSuite`, `DistributionSuite`, the `KeyGroupedPartitioning*` suites, `EnsureRequirementsSuite`, `GroupPartitionsExecSuite`, `ProjectedOrderingAndPartitioningSuite`, `PlannerSuite`, `WriteDistributionAndOrderingSuite`, `BucketedReadWithoutHiveSupportSuite`, the plan stability suites, the `*JoinSuite` suites and `AdaptiveQueryExecSuite` on branch-4.3, 1857 tests in all, plus `dev/lint-scala`. ### What changes were proposed in this pull request? A storage-partitioned join now pairs a partition transform whose argument is not a bare column, i.e. is an expression or a struct field, only with a transform over the same argument. It also stops dropping such an argument where it used to: - `KeyedShuffleSpec.isExpressionCompatible` pairs a transform with an argument that is not an `Attribute` only with a transform over the same argument shape, i.e. the argument with its one column replaced by a placeholder. So `bucket(4, b + 1)` pairs with `bucket(4, c + 1)` on `b = c`, but not with `bucket(4, x)`. Over the same shape the two compare as two transforms of columns do. Under `allowCompatibleTransforms`, `bucket(8, b + 1)` reduces onto `bucket(4, c + 1)`, since a reduce maps the partition keys and evaluates no argument. When an earlier join reduced the keys, two such transforms pair only when they were reduced through the same pairing. An identity side is never reduced onto such a transform. `TransformExpression` keeps its comparisons. It loses `withReference`, whose last callers now rebuild the transform with `copy`. - `KeyedShuffleSpec.canCreatePartitioning` asks the same question, so such a layout is never the one other children are shuffled onto. `createPartitioning` replaces a transform's argument with the other child's cluster key, which drops whatever surrounds the key. It and the identity arm of `reducersBothWays` now rebuild a transform through one helper that asserts the argument is a column. - `ShuffleSpecCollection.canCreatePartitioning` is now true when any member can, and `EnsureRequirements.pickCoPartitionTarget` offers and ranks only the members that can. So one member that cannot serve no longer rules out a usable sibling. That also covers a bare expression such as `b + 1`, which master already refused. - `GroupPartitionsExec` rebuilds a reduced expression over each member's own argument when it reports the members of a collection. It used to re-target the column only. - `KeyedShuffleSpec.keyPositions` maps an expression with more than one reference, e.g. `bucket(4, b + c)`, to no position instead of failing an assertion (SPARK-59901). - `EnsureRequirements.createKeyedShuffleSpecs` counts only an expression over one column towards `spark.sql.requireAllClusterKeysForCoPartition`. An expression over two columns maps to no position, so a projection would drop it together with its columns. ### Why are the changes needed? With `spark.sql.sources.v2.bucketing.shuffle.enabled` on, this query returns 0 rows instead of 7. `t1(id)` and `t3(x)` are partitioned by `bucket(4, ...)`, `plain(b)` is not partitioned, and each holds the values 0 to 7: ```sql SELECT t1.id, p.b, t3.x FROM t1 JOIN plain p ON t1.id = p.b + 1 JOIN t3 ON p.b = t3.x ``` 1. The first join shuffles `plain` onto `t1`'s partitioning. `KeyedShuffleSpec.createPartitioning` builds that partitioning from `plain`'s join key, so it is `bucket(4, b + 1)`. This join is correct. 2. The second join clusters `plain` on the bare `b`. `KeyedShuffleSpec.keyPositions` maps a partition expression to a cluster key through its reference, so `bucket(4, b + 1)` counts as a function of `b`, the counterpart of `x`. 3. `isSameFunction` compares only the function name and the bucket count. So `bucket(4, b + 1)` is the same as `t3`'s `bucket(4, x)`, and the join pairs the partitions as they stand. A row with `b = x` sits in bucket `(b + 1) % 4` on one side and in `x % 4` on the other. The comparison ignores which column an argument is, and relies on `keyPositions` to pair the columns up. That is sound only 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`. A scan never reports a transform of an expression or of a struct field, but two planner paths build one. A one-side shuffle does, under `spark.sql.sources.v2.bucketing.shuffle.enabled`. An inner broadcast hash join does too, with every conf at its default. `BroadcastHashJoinExec.expandOutputPartitioning` also reports the streamed side's layout over the build side's join key. The same problem shows up in more shapes, all measured on master: - **A broadcast join.** With every conf at its default, `SELECT /*+ BROADCAST(p) */ ... FROM t1 JOIN plain p ON t1.id = p.b + 1 JOIN t3 ON p.b = t3.x` returns 0 of 7 rows the same way, through the expanded `bucket(4, b + 1)` layout. - **The reducer path.** Under `spark.sql.sources.v2.bucketing.allowCompatibleTransforms.enabled`, a `bucket(8, x)` third table loses the rows the same way. - **A struct field.** A join key `p.s.a` gives `bucket(4, s.a)`, which `keyPositions` pairs through `s`. Joining two such sides on `s` pairs `bucket(4, s.a)` with `bucket(4, s.b)` and returns 0 of 8 rows. Shuffling another side onto such a layout builds `bucket(4, s)`, which fails on executors with a `ClassCastException`. - **Shuffling onto the layout.** In `plain p JOIN t1 ON p.b + 1 = t1.id JOIN xy q ON t1.id = q.x AND p.b = q.y`, the second join shuffles `xy` onto the first join's `bucket(4, b + 1)` layout. That builds `bucket(4, y)` for `xy`, drops the `+ 1`, and returns 0 of 8 rows. - **After a reduce.** When a later join reduces the first join's output from `bucket(4)` onto `bucket(2)`, `GroupPartitionsExec` reports the `bucket(4, b + 1)` member as `bucket(2, b)`. A join on `b` then pairs it with a `bucket(2, x)` scan, or shuffles another side onto it, and returns 0 of 7 rows. The same happens to the bare `b + 1` member of a side shuffled onto an identity-partitioned table. The struct form fails with the `ClassCastException`. - **An identity side reduced onto it.** Under `allowCompatibleTransforms`, a later join can reduce an identity-partitioned side onto `bucket(4, b + 1)`. That evaluates `bucket(4, id + 1)` on the side's partition keys while planning. The keys come out right, but under ANSI a key of `Long.MaxValue` fails the query with `ARITHMETIC_OVERFLOW`, although the query never computes `id + 1`. - **A GROUP BY after a reduce.** With `allowCompatibleTransforms`, `SELECT /*+ BROADCAST(p) */ p.b, max(p.c) FROM ident i JOIN plain2 p ON i.id = p.b + p.c JOIN bucket2 t2 ON i.id = t2.id GROUP BY p.b` returns 4 groups instead of 2. The reduce reports the `p.b + p.c` member as `bucket(2, p.b)`, which serves the `GROUP BY` as it stands. - **A key over two columns.** A join key such as `p.b + p.c` gives `bucket(4, b + c)`, or the bare `b + c` over an identity-partitioned side. Any spec over it fails planning with an `AssertionError` in `keyPositions` (SPARK-59901). AQE validates the plan in `OptimizeSkewedJoin`, so a single join is enough. An `ORDER BY` over such a partition expression had a related problem in another code path. That was SPARK-59905, fixed in #59189. The one-side shuffle of transform expressions came with SPARK-48012, in 4.0.0. So did the broadcast expansion of a keyed layout, since SPARK-49205 made `KeyGroupedPartitioning` an expression. ### Does this PR introduce _any_ user-facing change? Yes, it fixes the wrong results, the crash and the planning failures above. Such a query now shuffles the side instead of pairing it or shuffling onto it. That costs shuffles in two shapes whose plan was already correct: - A widening cast counts as part of the argument's shape. When `t1.id` is `BIGINT` and `p.b` is `INT`, the first join reports `bucket(4, cast(b as bigint))`. A later join `p.b = t3.x` with an `INT` bucketed `t3` then shuffles the first join's output. Before, a connector whose bucket function has one canonical name for both types paired it as it stood. This happens with every conf at its default, through a broadcast join. Two sides that both carry the cast still pair. - Under `allowCompatibleTransforms`, an identity-partitioned side is no longer reduced onto such a transform, so it is shuffled. One plan gets better. A collection with a bare expression member, e.g. the `b + 1` of a side shuffled onto an identity-partitioned table, used to be refused as a layout as a whole. Now its other members can serve, so `ident i JOIN plain p ON i.id = p.b + 1 JOIN xy q ON i.id = q.x AND p.b = q.y` takes 2 shuffles instead of 3. This PR keeps the fix small, since it has to reach the maintenance branches. SPARK-59900 is the follow-up for the rest: shuffling another side onto such a layout, and reducing an identity side onto it. ### How was this patch tested? New tests. All but two of them fail on master: - `ShuffleSpecSuite`, "SPARK-59887: a spec over a transform of an expression pairs only with the same shape". For an expression and a struct field argument, it covers the same shape, another shape and a column, both ways and with itself, the reduce in both directions, the two sides of one reduce, two sides reduced through different pairings, an identity side, `canCreatePartitioning`, and a collection with a usable sibling. - `ShuffleSpecSuite`, "SPARK-59901: an expression over two columns maps to no cluster key". - `KeyGroupedPartitioningSuite`, "SPARK-59887: a side shuffled onto a transform of an expression is not paired as it is". It runs the query above against a `bucket(4)` third table, and with `allowCompatibleTransforms` against a `bucket(8)` one that the join reduces, and asserts two shuffles. - `KeyGroupedPartitioningSuite`, "SPARK-59887: a side is not shuffled onto a transform of an expression". It runs both join orders and asserts that `xy` is shuffled onto `bucket(4, x)` and never onto a bucket of `y`. - `KeyGroupedPartitioningSuite`, "SPARK-59887: a broadcast join's layout over a join key is not paired as it is". It runs with every conf at its default and checks that the first join is a broadcast join. - `KeyGroupedPartitioningSuite`, "SPARK-59887: a layout over a transform of a struct field is not paired nor shuffled onto". Master returns 0 of 8 rows. - `KeyGroupedPartitioningSuite`, "SPARK-59887: a reduce keeps what each member's keys are computed from". It covers a bucketed and an identity-partitioned first table, a bucketed and an unpartitioned third table, and the struct form. Master returns 0 of 7 rows. - `KeyGroupedPartitioningSuite`, "SPARK-59887: an identity side is not reduced onto a transform of an expression". It runs with `p.b + 1` over a `Long.MaxValue` key and with `-p.b` over a `Long.MinValue` key. Master fails with the `ARITHMETIC_OVERFLOW` for both. - `KeyGroupedPartitioningSuite`, "SPARK-59887: a reduce keeps a member over two columns". It runs the `GROUP BY` above. Master returns 4 groups instead of 2. - `KeyGroupedPartitioningSuite`, "SPARK-59901: a join key over two columns is not paired through one of them". Master fails with the `AssertionError`. - `KeyGroupedPartitioningSuite`, "SPARK-59901: a join key over two columns does not cover its columns". It joins two broadcast joins that report `[b + c, d]` on `(b, c, d)` under `allowKeysSubsetOfPartitionKeys`, and asserts that the join is not paired on `d` alone. Master fails with the `AssertionError`. - `KeyGroupedPartitioningSuite`, "SPARK-59887: two layouts over the same shape of a join key pair as two columns do". It runs the cast case above with every conf at its default. It also runs two one-side shuffles onto `bucket(4, b + 1)` and onto `bucket(4)` or `bucket(8)` of `q.b + 1`, where `allowCompatibleTransforms` reduces the `bucket(8)` one. It asserts that the join adds no shuffle. It passes on master, which pairs and reduces these as well. - `KeyGroupedPartitioningSuite`, "SPARK-59887: the same shape reduced together pairs only through the same pairing". Two legs each shuffle onto `bucket(12, b + 1)`. The first leg reduces with `bucket(8)`, so its keys are buckets of 4. The second leg reduces with `bucket(8)` too, with `bucket(18)`, or not at all, so its keys are buckets of 4, 6 or 12. It asserts 2 shuffles for the first and 4 for the other two. Master returns 0 of 23 rows for the third. - `KeyGroupedPartitioningSuite`, "SPARK-59887: a member that cannot serve does not rank its collection". It pins the ranking: `q` keeps its layout with 3 partitions, as on master, where the collection is refused as a whole. Without the member filter, the collection would rank first with 5 partitions, and `q` would be shuffled onto 2. The SPARK-59905 test from #59189 now also runs its `p.b + p.c` query with AQE on. Master fails that run with the `AssertionError`. The three expression tests each run with the key `p.b + 1` and with `-p.b`. Master returns 0 of 8 rows for `p.b + 1` and 4 of 8 for `-p.b`. The `-p.b` key is the shape that also fails on 4.2, 4.1 and 4.0. On those branches `b + 1` already works, since its literal leaf blocks the pairing. Each change was removed on its own: - without the guard in `isExpressionCompatible`, the one-side shuffle and broadcast pairing tests, the struct field test and the reduce test return wrong rows, and the identity tests fail on the assertion in `rebuiltOver`; - without the shape pairing, the same-shape and same-pairing tests fail on the shuffle count; - with a reduce of such a pair refused, the same-shape test fails on the shuffle count; - with every pair of reduced keys over an expression refused, the same-pairing test fails on the shuffle count; - with the reduced marks ignored, or only checked on both sides, the same-pairing test returns 7 of 23 rows; - without the `canCreatePartitioning` clause, the one-side shuffle pairing, `xy`, struct field, reduce and same-pairing tests fail, all on the assertion in `rebuiltOver`; - with the `forall` back in `ShuffleSpecCollection.canCreatePartitioning`, the `xy` tests fail on the shuffle count; - without the member filter in `pickCoPartitionTarget`, 15 tests fail in `KeyGroupedPartitioningSuite` and `EnsureRequirementsSuite`. The filter now also does what the collection-level check before it did, so 8 of them are older tests. Of this PR's tests, 7 fail, 6 of them on the assertion in `rebuiltOver`. A member over the same shape pairs with itself, so only the filter keeps it from being the layout; - without the `GroupPartitionsExec` change, the reduce test returns 0 of 7 rows, and the two-column `GROUP BY` returns 4 groups; - without the `requireAllClusterKeysForCoPartition` change, the coverage test pairs on `d` alone; - with the `keyPositions` assertion back, both SPARK-59901 tests in `KeyGroupedPartitioningSuite` fail. The first `ShuffleSpecSuite` test also fails for each change to `isExpressionCompatible` above. The `xy` tests pass with the guard removed, since the `canCreatePartitioning` clause and the member filter keep `xy` off that layout. Also ran `TransformExpressionSuite`, `ShuffleSpecSuite`, `DistributionSuite`, the `KeyGroupedPartitioning*` suites, `WriteDistributionAndOrderingSuite`, `PlannerSuite`, `ProjectedOrderingAndPartitioningSuite`, `GroupPartitionsExecSuite`, `EnsureRequirementsSuite`, `BucketedReadWithoutHiveSupportSuite`, the plan stability suites, the `*JoinSuite` suites and `AdaptiveQueryExecSuite`, 1955 tests in all, plus `dev/lint-scala`. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Claude Opus 5.5) --- .../expressions/TransformExpression.scala | 10 - .../plans/physical/partitioning.scala | 160 +++++-- .../spark/sql/catalyst/ShuffleSpecSuite.scala | 107 ++++- .../datasources/v2/GroupPartitionsExec.scala | 17 +- .../exchange/EnsureRequirements.scala | 27 +- .../KeyGroupedPartitioningSuite.scala | 444 +++++++++++++++++- .../exchange/EnsureRequirementsSuite.scala | 64 +++ 7 files changed, 764 insertions(+), 65 deletions(-) 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)),