peter-toth commented on code in PR #59165:
URL: https://github.com/apache/spark/pull/59165#discussion_r4167940664
##########
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:
Thanks, measured, and you're right. Master pairs the two as they stand with
no shuffle, and the guard shuffled both sides. Rather than only documenting it,
[765c526](https://github.com/apache/spark/pull/59165/commits/765c526071fd7c35ba80b518564272a344e64036)
pairs a transform of an expression with a transform over the same argument
shape, i.e. the argument with its one column replaced by a placeholder. Your
query, and two one-side shuffles onto `bucket(4, b + 1)` and `bucket(4, c +
1)`, add no shuffle again ("two layouts over the same shape of a join key pair
as two columns do"). Over the same shape the two compare as two transforms of
columns do, so `bucket(8, b + 1)` also reduces onto `bucket(4, c + 1)`. A
reduce maps the partition keys and evaluates no argument, so the ANSI concern
does not come up. An identity side is still not reduced onto such a transform.
The description now lists the costs that remain, and which of them apply with
default confs.
##########
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:
Confirmed and fixed in
[765c526](https://github.com/apache/spark/pull/59165/commits/765c526071fd7c35ba80b518564272a344e64036).
`allClusterKeysCovered` now counts only expressions over a single column. Your
shape is the new test "SPARK-59901: a join key over two columns does not cover
its columns", which paired on `d` alone before.
##########
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:
Thanks, measured. Master returns 4 groups for your query. Added as "a reduce
keeps a member over two columns" in
[765c526](https://github.com/apache/spark/pull/59165/commits/765c526071fd7c35ba80b518564272a344e64036).
##########
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:
Added in
[765c526](https://github.com/apache/spark/pull/59165/commits/765c526071fd7c35ba80b518564272a344e64036).
The identity test now runs `-p.b` over a `Long.MinValue` key as well, and
master fails that one with the `ARITHMETIC_OVERFLOW` too.
##########
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:
Measured, thanks. With your query the join lands on 3 partitions with the
filter and on 2 without it.
[765c526](https://github.com/apache/spark/pull/59165/commits/765c526071fd7c35ba80b518564272a344e64036)
adds it as "a member that cannot serve does not rank its collection", in
`KeyGroupedPartitioningSuite` rather than `EnsureRequirementsSuite`, since the
shape needs a broadcast join over partitioned catalog tables. The description
says the filter decides the ranking. With the filter the plan is the one master
takes, since master refuses that collection as a whole. With the shape pairing,
the filter is not just defensive any more either. A member over the same shape
pairs with itself, so without the filter the `xy` tests fail too.
`createPartitioning` now rebuilds the template through a helper that asserts
its argument is a column, so a regression there fails planning instead of
returning wrong rows.
##########
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:
Done in
[765c526](https://github.com/apache/spark/pull/59165/commits/765c526071fd7c35ba80b518564272a344e64036).
The doc names both producers now, and says the projection happens under
`allowKeysSubsetOfPartitionKeys` for an unmarked layout where another position
maps. The `createShuffleSpec` comment mentions the two-column case too. The
`EnsureRequirements` ordering comment went away with #59189.
##########
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:
Done in
[765c526](https://github.com/apache/spark/pull/59165/commits/765c526071fd7c35ba80b518564272a344e64036).
The note now states the two conditions: a partition expression that
`keyPositions` maps to no cluster key, or one that is neither a column nor a
transform. A spec over a transform of an expression is reflexive again, through
the shape pairing.
##########
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:
Done in
[765c526](https://github.com/apache/spark/pull/59165/commits/765c526071fd7c35ba80b518564272a344e64036).
The `IdentityReducer` comment credits `keyPositions` with the single
reference, and step 3.3 mentions the argument condition.
--
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]
