dongjoon-hyun commented on code in PR #59165:
URL: https://github.com/apache/spark/pull/59165#discussion_r4159568724
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -1121,8 +1125,8 @@ 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.
+ // Every column a partition expression references counts as covered. A
scan reports one
Review Comment:
**A member over two columns passes this check, and
`KeyedPartitioning.createShuffleSpec` then projects it away.** This counts
every column that `p.b + p.c` references as covered, but `keyPositions` now
maps `p.b + p.c` to no position (partitioning.scala:1964-1968). So under
`spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled`,
`createShuffleSpec` drops that position (partitioning.scala:1015). For
instance, take two sides that each report `[p.b + p.c, p.d]`, e.g. the output
of a broadcast join `ON i.x = p.b + p.c AND i.d = p.d` over a table partitioned
by `identity(x), identity(d)`. A join of the two on `(b, c, d)` passes this
check with `{b, c, d}`, but both specs are projected to `[p.d]`, and the
storage-partitioned join pairs them on `d` alone. The rows are right, but the
join runs on as many partitions as there are distinct values of `d`, which is
the skew `spark.sql.requireAllClusterKeysForCoPartition` is there to prevent.
On the base the same plan failed th
e `keyPositions` assertion, so it is a newly reachable plan rather than a
regression. Could this count only the columns of expressions with a single
reference, or check the projected spec instead?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -2061,7 +2074,15 @@ case class KeyedShuffleSpec(
left: Expression,
right: Expression,
allowReduce: Boolean): Boolean = {
- if (TransformExpression.hasReducedKeys(left) ||
TransformExpression.hasReducedKeys(right)) {
+ val overExpression = Seq(left, right).exists {
Review Comment:
**The costs listed in the description now also reach plans that use only
default confs, through the broadcast expansion.** The description says an inner
broadcast hash join builds such a layout with every conf at its default, but
the costs still read as if they need
`spark.sql.sources.v2.bucketing.shuffle.enabled`. For example, take `orders`
and `returns` bucketed by `bucket(16, cust_id)` on a `BIGINT` column, and an
`INT` `customers.id` broadcast into both:
```sql
SELECT a.id FROM
(SELECT c1.id FROM orders o JOIN customers c1 ON o.cust_id = c1.id) a
JOIN
(SELECT c2.id FROM returns r JOIN customers c2 ON r.cust_id = c2.id) b
ON a.id = b.id
```
Each broadcast join also reports `bucket(16, cast(cN.id as bigint))`. The
base pairs the two as they stand, with no shuffle, and that is correct since
both sides have the same shape. With this guard, both fact-sized sides are
hash-shuffled, since keyed shuffles are off by default. The `t.id = d.k + 1`
shape behaves the same. It is the first cost in the list, but on the
maintenance branches it changes plans for users who set nothing. Could the
description say that these costs apply with default confs too?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -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. Its argument is a bare column,
since the pairing
+ // admits no other. The member's own argument takes its
place, since that need
+ // not be one. 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 {
Review Comment:
**This also fixes a `GROUP BY` or window over a member with two columns, but
no test covers that.** The base re-targeted a `p.b + p.c` member at its first
reference, `bucket(2, p.b)`, which satisfies `ClusteredDistribution([p.b])`, so
a `GROUP BY p.b` or `PARTITION BY p.b` above it skipped its shuffle. With
`spark.sql.sources.v2.bucketing.allowCompatibleTransforms.enabled` on, `plain2`
holding `(0, 1), (0, 2), (1, 1), (1, 2)`, and the `ident` and `bucket2` tables
of the new tests:
```scala
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)
```
The base returns `(0, 2), (1, 1), (0, 1), (1, 2)`, since the rows of `b = 0`
sit in two buckets, and this PR returns `(0, 2), (1, 2)`. The aggregate has to
keep `p.c`: with `count(*)`, column pruning drops the member before the second
join. The reduce test (KeyGroupedPartitioningSuite.scala:1647) covers
single-column members under a later join only, and the SPARK-59901 test
(KeyGroupedPartitioningSuite.scala:1707) has no reduce. Could one of them add
this shape, or its `row_number() OVER (PARTITION BY p.b ORDER BY p.c)` variant?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -367,12 +369,14 @@ 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
+ // Choose all the specs that can be used to shuffle other children,
keeping only the members
+ // that can serve as the layout. Any other member would build a
partitioning its own child is
+ // not laid out on.
val candidateSpecs = children.zip(specs).collect {
case (child, Some(spec)) if spec.canCreatePartitioning &&
(!shouldConsiderMinParallelism ||
child.outputPartitioning.numPartitions >=
conf.defaultNumShufflePartitions) =>
- child -> spec
+ child -> spec.flatten.filter(_.canCreatePartitioning)
Review Comment:
**This filter is not only defensive: it can change which child the layout
comes from.** The description says that removing it changes no test, since the
guard already keeps a refused member from pairing with its own side. That holds
for the pairing, but the ranking at L391 now takes the max over the serving
members only, and with
`spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled` the
members of one collection can project to different counts. For example, with
that conf and `spark.sql.sources.v2.bucketing.shuffle.enabled` on, take `t1`
partitioned by `identity(a), identity(c)` with keys `(1, 1), (1, 2), (1, 3),
(2, 4), (2, 5)`, and `SELECT /*+ BROADCAST(p) */ * FROM t1 JOIN p ON t1.c = p.y
+ 1 JOIN q ON t1.a = q.m AND p.y = q.n`, where `q` reports 3 partitions without
an exchange. The broadcast join reports `[a, c]` and `[a, y + 1]`. The first
projects to `[a]`, 2 partitions, and can serve. The second keeps 5 and cannot.
With the filter, `q` wins with 3, as on
the base. Without it, `t1` ranks at 5 and wins, and `q` is shuffled onto 2
partitions. Both plans are correct, so only the parallelism differs. Could the
description say so, and could a test in EnsureRequirementsSuite pin the ranking?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1588,15 +1584,21 @@ 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 is neither a
column nor a transform
Review Comment:
Minor: following up on R1-10, this list reads as complete, but it misses
specs over a column, or a transform of a column, whose `keyPositions` entry is
empty. `years(ts)` clustered on `years(ts)` maps to nothing, since
`keyPositions` looks up `ts` (EnsureRequirementsSuite, "SPARK-59289: a
partitioning is never projected onto no position"). So does the unprojected
spec of a marked narrowing projection (ShuffleSpecSuite, "createShuffleSpec: a
marked narrowing projection yields an unusable spec"). Both fail
`left.intersect(right).nonEmpty` in `areKeysCompatible` against themselves,
while the last sentence still holds for them. Could the note state the two
conditions, a partition expression that maps to no cluster key in
`keyPositions`, or one that is neither a column nor a transform of columns,
instead of listing shapes?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1941,8 +1943,16 @@ 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
Review Comment:
Minor: this paragraph names `createPartitioning` as the only producer, and
its last sentence holds only in a narrow case. Since 5375cc3, the description
also names the broadcast expansion, which builds such an expression with
default confs. And `createShuffleSpec` projects the position away only under
`spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled`, for an
unmarked layout where another position maps. It returns a marked layout
unprojected at L1030-L1031, and every `createPartitioning` output is marked
(L2270-L2272). It returns the single `bucket(4, b + c)` of this example
unprojected too, at L1039-L1040, and the comment at L1036-L1037 still
attributes that only to "a clustering that names the partition expression
itself". EnsureRequirements.scala:143-145 likewise mentions only "a side
shuffled onto another side's layout". Could these comments name both producers,
and say when the projection happens?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -1520,6 +1534,203 @@ 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 = {
+ createTable("plain", Array(Column.create("b", LongType)), Array.empty)
+ sql("INSERT INTO testcat.ns.plain VALUES " + (0 until 8).map(b =>
s"($b)").mkString(", "))
+ 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()
+ createTable("plain", Array(Column.create("b", LongType)), Array.empty)
+ sql("INSERT INTO testcat.ns.plain VALUES " + (0 until 8).map(b =>
s"($b)").mkString(", "))
+ 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)))
+ }
+ }
+
+ test("SPARK-59887: an identity side is not reduced onto a transform of an
expression") {
+ createBucketedIdTable("bucket4", 4, numIds = 8)
+ createTable("plain", Array(Column.create("b", LongType)), Array.empty)
+ sql("INSERT INTO testcat.ns.plain VALUES " + (0 until 7).map(b =>
s"($b)").mkString(", "))
+ createIdentityIdTable((0L until 8L) :+ Long.MaxValue)
+
+ // The first join shuffles `plain` onto `bucket(4, b + 1)`. Reducing
`ident` onto it would
+ // evaluate `id + 1` on `ident`'s partition keys while planning, and that
overflows for
+ // `Long.MaxValue` under ANSI. The query itself never computes it, since
no `b` matches that
+ // key.
+ withSQLConf(
+ SQLConf.ANSI_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ val df = sql(
+ """SELECT k.id, p.b, i.id FROM testcat.ns.bucket4 k
+ |JOIN testcat.ns.plain p ON k.id = p.b + 1
Review Comment:
**This test would pass on branch-4.2 without the fix.** It is the only
end-to-end test of the identity arm, and it runs with `p.b + 1` only. As the
description says, branch-4.2 never pairs `p.b + 1`, because its
`createKeyedShuffleSpec` compares the leaves of the partition expressions with
the clustering and sees the literal. But branch-4.2 has the identity arm too,
`IdentityReducer(t.withReference(a))`, so `-p.b` does pair with `ident` there,
and the planner evaluates `bucket(4, -id)` on the keys of `ident`. Under ANSI
that fails with `ARITHMETIC_OVERFLOW` for a key of `Long.MinValue`, although
the query never negates it. The two expression tests above run both keys for
the same reason. Could this test run `-p.b` as well, with `Long.MinValue` in
`ident`?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1892,8 +1894,8 @@ case class KeyReducer(reducer: Reducer[_, _],
reducedExpression: TransformExpres
* @param transform the partition transform, re-targeted at 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
(`KeyedShuffleSpec.isExpressionCompatible` admits no
Review Comment:
Nit: this credits `isExpressionCompatible` with the single column, but
`argumentsAreColumns` admits any number of columns. The single reference comes
from `keyPositions`, as the comment at L2203-L2205 says. Relatedly, step 3.3 of
`isCompatibleWith` (L1982-L1985) still says that "the same transform function,
or the two sides of one reduce" describe one key space, while the guard now
refuses both when an argument is not a column, as ShuffleSpecSuite.scala:807
and ShuffleSpecSuite.scala:815 assert. Could both comments mention the argument
condition?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1951,8 +1961,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)
Review Comment:
Minor: how about mapping by the argument here instead of by the reference?
`createPartitioning` replaces the argument of a transform, and
`GroupPartitionsExec` now rebuilds over it (GroupPartitionsExec.scala:611-615),
but this maps an expression by its single reference. If it looked up
`t.children.head` for a single-child transform, and the expression itself
otherwise, every shape the new tests refuse would map to no position:
`bucket(4, a + 1)` under `[a]`, `bucket(4, s.x)` under `[s]`, `a + b` under
`[a, b]`, and the reduced `bucket(2, p.b + 1)` under `[p.b]`. The overlap test
in `areKeysCompatible`, `canCreatePartitioning` and `createPartitioning` would
then agree by construction, without `argumentsAreColumns` or the `refs.size ==
1` branch. Today three shape checks have to be kept in step by hand, and they
already differ: `refs.size == 1` here, the leaf match after `overExpression` in
`isExpressionCompatible`, and `AttributeReference` or `argumentsAreColumns` in
`canCreatePart
itioning`. Unlike the SPARK-59900 idea, this evaluates nothing new, so the
ANSI concern does not arise. The identity arm at L2207 and L2211 would have to
switch to `t.copy(children = Seq(a))` at the same time. Otherwise, with
`spark.sql.requireAllClusterKeysForCoPartition` off, a join keyed on `p.b + 1`
would give `bucket(4, p.b + 1)` a position, and `withReference` would build
`bucket(4, id + 1)`. I understand this may fit SPARK-59900 better, so please
take it as input.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -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. Its argument is a bare column,
since the pairing
+ // admits no other. The member's own argument takes its
place, since that need
+ // not be one. 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)
Review Comment:
Nit: this replaces the use that `TransformExpression.withReference` was
added for in SPARK-59045, but its Scaladoc (TransformExpression.scala:127-131)
still relies on the single leaf of `KeyedPartitioning.supportsExpressions` and
warns about field paths. TransformExpressionSuite.scala:115-116 also still says
the marker survives the rewrites that "`GroupPartitionsExec` apply". The only
callers left are the identity arm (partitioning.scala:2207 and
partitioning.scala:2211), which sees a transform over a single bare column.
Could the identity arm use `t.copy(children = Seq(a))` like this line, so that
`withReference` and its test can go? Otherwise, could the two comments be
updated?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -139,8 +139,10 @@ case class EnsureRequirements(
// OrderedDistribution requires grouped KeyedPartitioning with
sorted keys
// according to the distribution's ordering.
val satisfyingKeyedPartitioning = resolution.fold(identity, _._1)
- // The single-column invariant in
KeyedPartitioning.supportsExpressions guarantees
- // one attribute per partition expression.
+ // This binds the ordering to the partition expressions'
references, which is right
Review Comment:
Nit: this rewrites one of four copies of the single-column comment, and
L1128-L1129 rewrite another. The two in `reorderJoinKeysRecursively`, at
L534-L535 and L541-L542, still say "The single-column invariant in
KeyedPartitioning.supportsExpressions guarantees one attribute per partition
expression". The SPARK-59901 test and the broadcast expansion now reach that
code with `p.b + p.c`, where `flatMap(_.references)` returns two columns for
one expression. `reorder` then finds no order, or only permutes key pairs, so
nothing breaks. Could those two be reworded the same way?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala:
##########
@@ -761,6 +761,84 @@ class ShuffleSpecSuite extends SparkFunSuite with
SQLHelper {
}
}
+ test("SPARK-59887: a spec over a transform of an expression pairs with
nothing") {
+ // 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 = b` it would pair
+ // `bucket(4, s.x)` with `bucket(4, b)`. A row with `a = b` sits in a
different bucket on each
+ // side.
+ val fn = new FakeBucket
+ val a = $"a".long
+ val b = $"b".long
+ val s = $"s".struct(new StructType().add("x", LongType))
+ // 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(argument: Expression, numBuckets: Int): TransformExpression =
Review Comment:
Nit: this reverses the argument order of `bucket(numBuckets, expr)` at L676,
of the connector `Expressions.bucket(numBuckets, ...)`, and of the comments of
this test, e.g. L782 says `bucket(8, argument)` while L784 calls
`bucket(argument, 8)`. Also, the `s = b` example at L766-L767 joins a struct
with a long, which the analyzer rejects. The shape that can happen is a join of
two structs, `s = t`, which pairs `bucket(4, s.a)` with `bucket(4, t.b)`. Could
the helper take the bucket count first, and the comment use `s = t`?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -1520,6 +1534,203 @@ 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 = {
+ createTable("plain", Array(Column.create("b", LongType)), Array.empty)
+ sql("INSERT INTO testcat.ns.plain VALUES " + (0 until 8).map(b =>
s"($b)").mkString(", "))
+ 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()
+ createTable("plain", Array(Column.create("b", LongType)), Array.empty)
+ sql("INSERT INTO testcat.ns.plain VALUES " + (0 until 8).map(b =>
s"($b)").mkString(", "))
+ 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)))
+ }
+ }
+
+ test("SPARK-59887: an identity side is not reduced onto a transform of an
expression") {
+ createBucketedIdTable("bucket4", 4, numIds = 8)
+ createTable("plain", Array(Column.create("b", LongType)), Array.empty)
+ sql("INSERT INTO testcat.ns.plain VALUES " + (0 until 7).map(b =>
s"($b)").mkString(", "))
Review Comment:
Nit: `plain` is created inline three times in the new tests (L1543-L1544,
L1650-L1651 and here), and only this copy inserts `0 until 7`, with no comment
saying why. The expected rows at L1703 stay the same with `0 until 8`, since
`p.b + 1 = 8` is not in `bucket4`. How about a `createPlainTable(bs: Seq[Long]
= 0L until 8L)` next to `createIdentityIdTable`? Also, the two new helpers sit
between `createBucketedIdTable` (L740) and its wrapper `createBucketedIdTables`
(L764), which were adjacent before. Could they move after that group?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -367,12 +369,14 @@ 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
+ // Choose all the specs that can be used to shuffle other children,
keeping only the members
+ // that can serve as the layout. Any other member would build a
partitioning its own child is
+ // not laid out on.
val candidateSpecs = children.zip(specs).collect {
case (child, Some(spec)) if spec.canCreatePartitioning &&
Review Comment:
Nit: `spec.canCreatePartitioning` here is the same as the member list below
being non-empty, since a collection now answers `exists` over its members. So
one decision has two sources, and the `.max` at L391 depends on that identity
in partitioning.scala. If `ShuffleSpecCollection.canCreatePartitioning`
changes, the `.max` throws on an empty list, which the SPARK-59256 test comment
warns about. Could this drop the guard and keep the children whose `members`
are non-empty instead? The comments at L385-L386 ("the best any member offers")
and L405 ("over all of its members") could then say serving members.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
