This is an automated email from the ASF dual-hosted git repository.
peter-toth pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new d8d72d2d3b18 [SPARK-59887][SPARK-59901][SQL] Fix wrong results when a
storage-partitioned join pairs a transform of a join-key expression with one of
a column
d8d72d2d3b18 is described below
commit d8d72d2d3b18d0357776c3cc06314482a25b3c7d
Author: Peter Toth <[email protected]>
AuthorDate: Mon Oct 5 17:14:35 2026 +0200
[SPARK-59887][SPARK-59901][SQL] Fix wrong results when a
storage-partitioned join pairs a transform of a join-key expression with one of
a column
### 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 ma [...]
- `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 [...]
- **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 three 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
both sides. Before, a connector whose bucket function has one canonical name
for both types paired the two sides as they 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.
- With `spark.sql.sources.v2.bucketing.shuffle.enabled`, a join on two keys
whose sides pair only through a transform of an expression gets one more
shuffle. Take two broadcast joins, `t1 JOIN /*+ BROADCAST(p) */ p ON t1.id =
p.b + 1` and the same over `t2` and `q`, joined on `l.id = r.c AND l.b = r.b`.
Each layout covers only one join key, so the storage-partitioned join declines
under `requireAllClusterKeysForCoPartition`. Master then keeps both sides,
paired on `bucket(4, b + 1)`. [...]
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)
Closes #59165 from peter-toth/SPARK-59887-spj-expression-key-identity.
Authored-by: Peter Toth <[email protected]>
Signed-off-by: Peter Toth <[email protected]>
---
.../catalyst/expressions/TransformExpression.scala | 10 -
.../sql/catalyst/plans/physical/partitioning.scala | 161 +++++++--
.../spark/sql/catalyst/ShuffleSpecSuite.scala | 107 +++++-
.../expressions/TransformExpressionSuite.scala | 4 -
.../datasources/v2/GroupPartitionsExec.scala | 15 +-
.../execution/exchange/EnsureRequirements.scala | 33 +-
.../connector/KeyGroupedPartitioningSuite.scala | 392 ++++++++++++++++++++-
7 files changed, 642 insertions(+), 80 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 191de9bf5062..bfa94914f4ae 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 bfc4d7ac8eab..f14410eb0ae5 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
@@ -26,6 +26,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
@@ -873,7 +874,7 @@ case class KeyedPartitioning(
* Not the same question as `KeyedShuffleSpec.keyPositions`, and the two
must not be merged. That
* one answers, per expression, which cluster key the expression is a
function *of*. That is
* what `KeyedShuffleSpec.createPartitioning` needs, since it rebuilds the
expression over the
- * other side's clustering with `te.copy(children = ...
clustering(positionSet.head))`. Handing it
+ * other side's clustering with `rebuiltOver(te,
clustering(positionSet.head))`. Handing it
* the second form above would rebuild a `years(ts)` clustered on
`years(ts)` as
* `years(years(ts))`.
* The two coincide exactly where every expression has a single reference
and no expression is
@@ -975,9 +976,9 @@ case class KeyedPartitioning(
* `keysSatisfy` may read the partition keys to decide whether a projection
would merge anything.
* `mayProjectToClusterKeys` holds the permission.
*
- * Given the single reference per expression that `supportsExpressions`
enforces, that case
- * coincides with `KeyedShuffleSpec.keyPositions.exists(_.nonEmpty)`, where
`createShuffleSpec`
- * takes its `joinKeyPositions` from. They cannot be merged into one
derivation, see
+ * Where every partition expression has a single reference, that case
coincides with
+ * `KeyedShuffleSpec.keyPositions.exists(_.nonEmpty)`, where
`createShuffleSpec` takes its
+ * `joinKeyPositions` from. They cannot be merged into one derivation, see
* `positionsCoveringClusterKeys`.
*/
private def keysCanSatisfy(required: Distribution): Boolean = {
@@ -1035,7 +1036,8 @@ case class KeyedPartitioning(
// 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. Return the unprojected spec, as
above.
+ // says yes and this finds nothing. It is reachable for an expression
over several columns
+ // too, e.g. `bucket(4, b + c)`. Return the unprojected spec, as above.
if (joinKeyPositions.isEmpty && expressions.nonEmpty) {
return result
}
@@ -1375,11 +1377,7 @@ case class PartitioningCollection(partitionings:
Seq[Partitioning])
// through here.
//
// Every admitted member stays, because `isCompatibleWith` answers for any
of them and the
- // collection cannot know which one the other side matched. The cost is
that
- // `ShuffleSpecCollection.canCreatePartitioning` is a `forall` over
members whose keyed half is
- // false without `v2BucketingShuffleEnabled`, so a groupable keyed member
beside a usable
- // non-keyed one would cost the collection its role as a shuffle template.
No operator is known
- // to report that mixture.
+ // collection cannot know which one the other side matched.
val filtered =
partitionings.filter(PartitioningCollection.maySatisfyAfterProjection(_,
distribution))
ShuffleSpecCollection(filtered.map(_.createShuffleSpec(distribution)))
@@ -1588,15 +1586,23 @@ sealed 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. A
- * [[ShuffleSpecCollection]] answers for the whole choice, since the planner
asks it of a child's
- * spec as a whole. Building the partitioning is
[[LeafShuffleSpec.createPartitioning]], and that
- * is always one member's job.
+ * [[ShuffleSpecCollection]] can when any of its members can. Building the
partitioning is
+ * [[LeafShuffleSpec.createPartitioning]], and that is always one member's
job, so the planner
+ * builds it from a member that can.
*/
def canCreatePartitioning: Boolean
@@ -1873,10 +1879,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)
@@ -1889,11 +1895,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))
@@ -1941,8 +1949,20 @@ case class KeyedShuffleSpec(
* `KeyedPartitioning.positionsCoveringClusterKeys` counts it. See that
method for why the two
* stay apart.
*
- * The assertion is what makes the reference reading sound. With one
reference per expression,
- * "some reference is a cluster key" and "every reference is" are the same
statement.
+ * Only an expression with a single reference maps to a position. With one
reference, "some
+ * reference is a cluster key" and "every reference is" are the same
statement, which is what
+ * makes the reference reading sound. 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]
@@ -1951,8 +1971,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
+ }
}
}
@@ -1967,9 +1990,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 &&
@@ -2061,7 +2085,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
@@ -2086,6 +2119,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.
@@ -2135,8 +2210,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.
@@ -2164,15 +2239,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)
}
@@ -2199,8 +2275,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.
@@ -2212,7 +2294,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 partitions, so it shares
their layout, with one
@@ -2289,7 +2372,7 @@ case class ShuffleSpecCollection(specs: Seq[ShuffleSpec])
extends ShuffleSpec {
}
override def canCreatePartitioning: Boolean =
- specs.forall(_.canCreatePartitioning)
+ specs.exists(_.canCreatePartitioning)
override def flatten: Seq[LeafShuffleSpec] = specs.flatMap(_.flatten)
}
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala
index 86eec4be3870..a73235118c57 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}
@@ -761,6 +761,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/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/TransformExpressionSuite.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/TransformExpressionSuite.scala
index db24614774f4..d333b4144161 100644
---
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/TransformExpressionSuite.scala
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/TransformExpressionSuite.scala
@@ -111,9 +111,5 @@ class TransformExpressionSuite extends SparkFunSuite {
assert(!left.hasSameReducedKeys(b12), "an unreduced side never shares one")
assert(!b12.hasSameReducedKeys(left))
assert(!b12.hasSameReducedKeys(b8), "nor do two unreduced ones")
-
- // The marker rides on the expression, so it survives the attribute
rewrites a projection and
- // `GroupPartitionsExec` apply to a reported partitioning.
- assert(left.hasSameReducedKeys(left.withReference(b)))
}
}
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 ea6adcd814dc..2cdc6b9d410b 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
@@ -570,8 +570,8 @@ private[sql] object GroupPartitionsExec {
// can only differ in `expressions`, since they share one `KeyLayout`
(enforced by
// `PartitioningCollection`). So the grouping is computed once, and the
layout it describes
// is shared by the members below.
- // 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, so the reduce marks it (see
`KeyedShuffleSpec.reducersBothWays`).
@@ -604,8 +604,15 @@ private[sql] object GroupPartitionsExec {
// `reduced` came from the one member
`checkKeyGroupCompatible` paired
// this side on, which need not be the member being
rewritten. The keys are
// reduced once, from the shared key rows, so `reduced`
describes them
- // whichever member this is, and only the key attribute is
re-targeted.
- reduced.withReference(expr.references.head)
+ // 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 b6a01a4ed73b..af186f4d0a62 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
@@ -377,25 +377,28 @@ case class EnsureRequirements(
val shouldConsiderMinParallelism = children.zip(specs).forall { case
(child, spec) =>
spec.forall(!_.canCreatePartitioning) ||
child.isInstanceOf[ShuffleExchangeLike]
}
- // Choose all the specs that can be used to shuffle other children
- val candidateSpecs = children.zip(specs).collect {
- case (child, Some(spec)) if spec.canCreatePartitioning &&
- (!shouldConsiderMinParallelism ||
- child.outputPartitioning.numPartitions >=
conf.defaultNumShufflePartitions) =>
- child -> spec
+ // Choose all the children whose spec can be used to shuffle other
children, each with the
+ // members that can serve as the layout. Any other member would build a
partitioning its own
+ // child is not laid out on, and it does not count towards the ranking
below either.
+ val candidateSpecs = children.zip(specs).flatMap {
+ case (child, Some(spec)) if !shouldConsiderMinParallelism ||
+ child.outputPartitioning.numPartitions >=
conf.defaultNumShufflePartitions =>
+ val members = spec.flatten.filter(_.canCreatePartitioning)
+ if (members.nonEmpty) Some(child -> members) else None
+ case _ => None
}
// Rank on two things at once. A child with no `ShuffleExchangeLike` node
comes first, since
// keeping it costs nothing. For instance, if we have:
// A: (No_Exchange, 100) <---> B: (Exchange, 120)
// it's better to pick A and change B to (Exchange, 100) instead of
picking B and insert a
// new shuffle for A. Then the best parallelism decides, and for a
collection that is the best
- // any member offers, since a collection has no count of its own.
+ // any serving member offers, since a collection has no count of its own.
//
// What the winner contributes is its members, the alternatives the layout
is picked from below.
// Empty when no child can serve as the layout.
- val bestMembers: Seq[LeafShuffleSpec] = candidateSpecs.maxByOption { case
(child, spec) =>
- (!child.isInstanceOf[ShuffleExchangeLike],
spec.flatten.map(_.numPartitions).max)
- }.toSeq.flatMap(_._2.flatten)
+ val bestMembers: Seq[LeafShuffleSpec] = candidateSpecs.maxByOption { case
(child, members) =>
+ (!child.isInstanceOf[ShuffleExchangeLike],
members.map(_.numPartitions).max)
+ }.toSeq.flatMap(_._2)
// A `ShuffleSpecCollection` answers `isCompatibleWith` if *any* of its
members does, so the
// winner alone does not say which member the sides agreed on. The
projection pushed into a
@@ -408,7 +411,7 @@ case class EnsureRequirements(
val pairings: Seq[(LeafShuffleSpec, Seq[Option[LeafShuffleSpec]])] =
bestMembers.map { member =>
member -> childMembers.map(_.find(member.isCompatibleWith))
}
- // Which children the winner reaches at all, over all of its members.
+ // Which children the winner reaches at all, over all of its serving
members.
val reached: Seq[Boolean] =
pairings.map(_._2).transpose.map(_.exists(_.isDefined))
// The member picked has to pair with every child the winner reaches. A
child it does not
// reach, and a child that is not in the decision, are shuffled whichever
member wins, so
@@ -1133,9 +1136,11 @@ case class EnsureRequirements(
// the skew of joining on keys that are coarser than the join keys. Key
order and duplicated
// cluster keys don't matter.
def allClusterKeysCovered: Boolean =
- // The single-column invariant in
KeyedPartitioning.supportsExpressions guarantees one
- // attribute per partition expression.
-
distribution.allClusterKeysAmong(partitioning.expressions.flatMap(_.references))
+ // 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 its columns are not covered.
+ distribution.allClusterKeysAmong(
+ partitioning.expressions.filter(_.references.size ==
1).flatMap(_.references))
// The coverage requirement is a comparison of expressions, while
`keysMaySatisfy` can end in
// a projection of the partition keys, so the cheap question is asked
first. The requirement
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 03b493b96727..ea1fafd94cfe 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
@@ -48,7 +48,7 @@ import org.apache.spark.sql.execution.{
import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec,
HashAggregateExec, SortAggregateExec}
import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec,
DataSourceV2ScanRelation, GroupPartitionsExec}
import org.apache.spark.sql.execution.exchange.{EnsureRequirements,
ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike,
ValidateRequirements}
-import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec,
ShuffledJoin, SortMergeJoinExec}
+import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec,
ShuffledHashJoinExec, ShuffledJoin, SortMergeJoinExec}
import org.apache.spark.sql.execution.metric.SQLMetricsTestUtils
import org.apache.spark.sql.execution.ui.SparkPlanGraphNode
import org.apache.spark.sql.execution.window.{Final, Partial,
WindowGroupLimitExec, WindowGroupLimitMode}
@@ -750,6 +750,33 @@ class KeyGroupedPartitioningSuite
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<third>` on `id`, in that order.
*/
private def threeWayBucketJoinDF(third: Int): DataFrame =
sql("SELECT b12.id FROM testcat.ns.bucket12 b12 " +
@@ -1529,6 +1556,357 @@ class KeyGroupedPartitioningSuite
}
}
+ // 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. AQE
+ // validates the first plan. 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)
+ }
+ }
+
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
@@ -5492,21 +5870,19 @@ class KeyGroupedPartitioningSuite
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(
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
