peter-toth commented on code in PR #59165:
URL: https://github.com/apache/spark/pull/59165#discussion_r4167942387
##########
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:
#59189 rewords both
([comment](https://github.com/apache/spark/pull/59189#discussion_r4158266464)),
and
[765c526](https://github.com/apache/spark/pull/59165/commits/765c526071fd7c35ba80b518564272a344e64036)
is rebased onto it.
##########
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:
Done in
[765c526](https://github.com/apache/spark/pull/59165/commits/765c526071fd7c35ba80b518564272a344e64036).
The identity arm rebuilds the transform with `copy`, through a helper it
shares with `createPartitioning`. `withReference` is gone together with its
`TransformExpressionSuite` assertion.
##########
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:
Thanks, that's a neat way to make the checks agree by construction, and it
evaluates nothing new either. I'll take it into SPARK-59900, since it also
changes what a position means for the projection in `createShuffleSpec`.
##########
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:
Done in
[765c526](https://github.com/apache/spark/pull/59165/commits/765c526071fd7c35ba80b518564272a344e64036).
The candidates are the children with a serving member, and the comments say
serving members.
##########
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:
Done in
[765c526](https://github.com/apache/spark/pull/59165/commits/765c526071fd7c35ba80b518564272a344e64036).
The helper takes the bucket count first, and the comment uses `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:
Done in
[765c526](https://github.com/apache/spark/pull/59165/commits/765c526071fd7c35ba80b518564272a344e64036).
`createPlainTable` takes the values, defaults to `0 until 8`, and sits after
`createBucketedIdTables` with the other helpers.
--
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]
