dongjoon-hyun commented on code in PR #59165:
URL: https://github.com/apache/spark/pull/59165#discussion_r4148432829
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -77,10 +87,21 @@ case class TransformExpression(
* It compares the transforms only. A caller that compares partition keys
has to consult
* `reducedWith` as well, since a reduced key space is not the one its
transform names.
*
+ * Only a transform whose arguments are all bare columns can be the same as
another. The
+ * comparison ignores which column an argument is, and
`KeyedShuffleSpec.keyPositions` pairs the
+ * two sides up by the column each expression references. That pairing is
sound only when the
+ * argument is the column itself. `KeyedShuffleSpec.createPartitioning`
builds a transform over
+ * the other side's join key, which can be an expression or a struct field.
`bucket(4, b + 1)` is
+ * a function of `b`, but not the same function of it as `bucket(4, x)` is
of `x`, so a join on
+ * `b = x` must not pair the two, and neither may a join on `s = t` pair
`bucket(4, s.a)` with
+ * `bucket(4, t.b)`. Such a transform is not the same as any transform,
itself included, which
+ * costs a shuffle rather than a match.
+ *
* @param other the transform expression to compare to
* @return true if this and `other` has the same semantics w.r.t to
transform, false otherwise.
*/
- def isSameFunction(other: TransformExpression): Boolean = functionId ==
other.functionId
+ def isSameFunction(other: TransformExpression): Boolean =
+ argumentsAreAttributes && other.argumentsAreAttributes && functionId ==
other.functionId
Review Comment:
**The new guard is bypassed once a later join reduces the keys, and the
wrong result comes back.** `GroupPartitionsExec` rewrites every member of the
reduced child's partitioning with the reduced expression re-targeted at the
member's first reference:
https://github.com/apache/spark/blob/14141b9e262ccc61b5ba8b721874a288965c7b9a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala#L603-L608
So a `bucket(4, b + 1)` member, which the one-side shuffle created, is
reported as `bucket(2, b)` after a join reduces `bucket(4)` onto `bucket(2)`.
Its argument is a bare column, so `isSameFunction` pairs it and
`canCreatePartitioning` accepts it, while its rows still sit in `(b + 1) % 2`.
The following returns 0 of 7 rows on both this PR and its base commit:
```scala
Seq("bucket4" -> 4, "bucket2a" -> 2, "bucket2b" -> 2).foreach { case (name,
n) =>
createBucketedIdTable(name, n, numIds = 8)
}
createTable("plain", Array(Column.create("b", LongType)), Array.empty)
sql("INSERT INTO testcat.ns.plain VALUES " + (0 until 8).map(b =>
s"($b)").mkString(", "))
withSQLConf(
SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
val df = sql(
"""SELECT b4.id, p.b, t2.id, t3.id FROM testcat.ns.bucket4 b4
|JOIN testcat.ns.plain p ON b4.id = p.b + 1
|JOIN testcat.ns.bucket2a t2 ON b4.id = t2.id
|JOIN testcat.ns.bucket2b t3 ON p.b = t3.id""".stripMargin)
checkAnswer(df, (0 until 7).map(b => Row(b + 1L, b.toLong, b + 1L,
b.toLong)))
}
```
The second join reduces the left side's `bucket(4, id)` onto `bucket(2)`,
and the third join pairs the re-targeted `bucket(2, b)` with `bucket(2, id)`
without any exchange. The same happens in three more shapes:
- When the third table is unpartitioned, it is shuffled onto `bucket(2, y)`,
and the query returns 0 of 7 rows.
- When the first table is identity-partitioned, the identity reduce turns
the bare `b + 1` member into `bucket(2, b)`, and the query returns 0 of 7 rows.
- In the struct form, `bucket(4, s.a)` becomes `bucket(2, s)`, and another
side shuffled onto it fails with the same `ClassCastException` as the
struct-field tests of this PR.
Could `GroupPartitionsExec` re-target by the member's own argument instead?
That means replacing the argument of `reduced` with the member's argument, i.e.
the child of a transform, or the member itself for a bare expression. The
member is then reported as `bucket(2, b + 1)`, which describes its keys and
which the new guard refuses.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -2199,8 +2203,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) &&
Review Comment:
**A two-column join-key expression fails planning in `keyPositions`, before
the new clause below is reached.** `createPartitioning` builds the partition
expression over the other child's join key, which can reference two columns,
e.g. `bucket(4, b + c)`, or the bare `b + c` for an identity-partitioned side.
Any spec over it then fails the assertion at L1958. This single join fails on
both this PR and its base commit:
```scala
createTable("ident", Array(Column.create("id", LongType)),
Array(identity("id")))
sql("INSERT INTO testcat.ns.ident VALUES " + (0 until 8).map(i =>
s"($i)").mkString(", "))
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)")
withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
sql("SELECT * FROM testcat.ns.ident i JOIN testcat.ns.plain2 p ON i.id =
p.b + p.c")
.collect()
}
```
```
[INTERNAL_ERROR] The Spark SQL phase planning failed with an internal error.
...
Cause: java.lang.AssertionError: assertion failed: Expected exactly one
child from
(b#10L + c#11L), but found 2
```
It is reached through AQE's `OptimizeSkewedJoin`, which validates the plan
with `ValidateRequirements`, so it needs nothing but
`spark.sql.sources.v2.bucketing.shuffle.enabled`. With a bucketed side,
`bucket4 b4 JOIN plain2 p ON b4.id = p.b + p.c JOIN xy q ON p.b = q.x AND p.c =
q.y` fails the same way in `EnsureRequirements.pickCoPartitionTarget`, through
this line.
I know this predates this PR, but it is the same `createPartitioning` shape.
Moving the new clause ahead of `keyPositions` fixes only the second path. So
`createPartitioning`, or `EnsureRequirements` when it picks the target, would
need to refuse a cluster key with more than one reference. Alternatively,
`keyPositions` could return an empty position set instead of asserting.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -2199,8 +2203,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 => t.argumentsAreAttributes
Review Comment:
**One refused member disables the whole collection, so the cost is larger
than the PR description says.** `ShuffleSpecCollection.canCreatePartitioning`
is a `forall`, so this clause also rules out the usable `bucket(4, id)` member
of the same join output. The keyed layout is then lost for every operator above
it too. I measured these with `spark.sql.sources.v2.bucketing.shuffle.enabled`
on, and both return correct results:
- The example in the PR description, `t1 JOIN plain p ON t1.id = p.b + 1
JOIN xy q ON t1.id = q.x AND p.b = q.y`, followed by `GROUP BY t1.id`, takes 4
shuffles instead of 2. `xy` and the first join's output are hash-shuffled on
`(x, y)` and `(id, b)`, and the aggregate then needs another exchange on `id`.
- Two one-side-shuffled legs, `(t1 JOIN p ON t1.id = p.b + 1) JOIN (t2 JOIN
q ON t2.id = q.c + 1) ON t1.id = t2.id AND p.b = q.c`, take 4 shuffles instead
of 2. On the base commit, the outer join pairs the two legs as they stand.
Could the refusal be per member instead? For instance,
`ShuffleSpecCollection.canCreatePartitioning` could be an `exists`, and
`EnsureRequirements.pickCoPartitionTarget` could keep only the winner's members
whose own `canCreatePartitioning` holds when it builds `bestMembers`. Then
`bucket(4, id)` stays usable, and `bucket(4, b + 1)` is never used as a
template. The change is as local as this clause, so it seems safe to backport
too.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1943,6 +1943,10 @@ case class KeyedShuffleSpec(
*
* 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.
+ *
+ * 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. Telling those
two apart is
Review Comment:
**The same reference-based reading also breaks `ORDER BY`.** When
`EnsureRequirements` checks a keyed child against an `OrderedDistribution`, it
binds the required ordering to `expressions.flatMap(_.references)`:
https://github.com/apache/spark/blob/14141b9e262ccc61b5ba8b721874a288965c7b9a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala#L142-L145
For a bare expression member that `createPartitioning` built over an
identity side, e.g. `100 - b`, the key rows already hold `100 - b`. The
ordering then computes `100 - (100 - b)` and sorts the partitions the wrong
way. The global sort is removed, and the rows come back in reverse order on
both this PR and its base commit:
```scala
// `ident` is identity-partitioned on `id` and holds 94..100, and `plain(b)`
holds 0..6.
withSQLConf(
SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
SQLConf.V2_BUCKETING_SORTING_ENABLED.key -> "true") {
val df = sql(
"""SELECT p.b FROM testcat.ns.ident i JOIN testcat.ns.plain p ON i.id =
100 - p.b
|ORDER BY 100 - p.b""".stripMargin)
// Returns 0, 1, ..., 6 instead of 6, 5, ..., 0.
assert(df.collect().map(_.getLong(0)).toSeq == (0 to
6).reverse.map(_.toLong))
}
```
This isn't from this PR either. But it is one more consumer that relies on
the argument being the column itself, like the one this paragraph describes, so
it may be worth covering here or in a separate JIRA.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -65,6 +65,16 @@ case class TransformExpression(
lazy val functionId: TransformFunctionId =
TransformFunctionId(function.canonicalName(), numBucketsOpt)
+ /**
+ * Whether every argument is a bare column, an [[Attribute]]. See
[[isSameFunction]] for why only
+ * such a transform can be compared with another one. A literal argument is
a parameter of the
+ * transform rather than a column, so it does not count.
+ */
+ private[sql] lazy val argumentsAreAttributes: Boolean = children.forall {
+ case _: Literal => true
+ case c => c.isInstanceOf[Attribute]
Review Comment:
**This also refuses an implicit widening `Cast` on the join key.** When
`t1.id` is `BIGINT` and `p.b` is `INT`, type coercion casts only the `INT`
side, so the first join shuffles `p` onto `bucket(4, cast(b as bigint))`. A
second join `p.b = t3.x`, with `t3` bucketed on an `INT` column `x`, was paired
as it stood before this PR. For a connector whose bucket function has one
canonical name for `INT` and `BIGINT`, that was a correct plan, and with this
PR it costs one more shuffle of the first join's output. The PR description
doesn't list this shape among the costs. Connectors with type-specific
canonical names are not affected, since their functions never compared equal.
If this matters, a widening cast of an attribute that `Cast.canUpCast` allows
could count as the column.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -1520,6 +1520,113 @@ 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, which
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`, so pairing the two 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`,
+ // so a row with `b = y` lands in a different bucket on each side.
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ val df = sql(
+ s"""SELECT k.id, p.b, q.x FROM testcat.ns.plain p
Review Comment:
**This test fails on the base commit only because `plain` is the left side
of the first join.** Both members of the first join's output have 4 partitions,
and `pickCoPartitionTarget` breaks the tie with `maxByOption`, which returns
the first member, i.e. `plain`'s. With the keyed table on the left, the base
commit already picks `bucket(4, id)` and returns the right answer. I checked
the `p.b + 1` case on the base with `bucket4 b4 JOIN plain p ON b4.id = p.b + 1
JOIN xy q ON b4.id = q.x AND p.b = q.y`: `xy` is shuffled onto `bucket(4, x)`,
and the answer is correct. Since this test checks only the answer, a change in
the member order or in how members are ranked would let it pass even with both
clauses reverted. How about asserting the plan too, e.g. that `xy` is not laid
out on a bucket of `y`, or the number of shuffles?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/TransformExpressionSuite.scala:
##########
@@ -116,4 +135,32 @@ class TransformExpressionSuite extends SparkFunSuite {
// `GroupPartitionsExec` apply to a reported partitioning.
assert(left.hasSameReducedKeys(left.withReference(b)))
}
+
+ test("SPARK-59887: a transform of an expression is not the same as any
transform") {
+ // A join pairs its two sides up by the column each partition expression
references, so over
+ // `a = b` it would pair `bucket(4, a + 1)` with `bucket(4, b)`. A row
with `a = b` sits in a
+ // different bucket on each side.
+ val fn = new NamedFunction("test.bucket")
+ val overExpression = bucket(fn, Add(a, Literal(1)))
+ assert(bucket(fn, a).isSameFunction(bucket(fn, b)), "two columns are the
same function")
+ assert(!overExpression.isSameFunction(bucket(fn, b)))
+ assert(!bucket(fn, b).isSameFunction(overExpression))
+ assert(!overExpression.isSameFunction(overExpression), "not even as
itself")
+
+ // A struct field is not the column the pairing uses either: that is the
struct.
Review Comment:
**`reducers` and `isCompatible` are not tested for a struct-field
argument.** The PR description says this test covers `isSameFunction`,
`reducers` in both directions and `isCompatible` for transforms over `a + 1`
and over a struct field. For the struct field, only `isSameFunction` is
checked, while `reducers` and `isCompatible` are checked with `Add(a,
Literal(1))` only. Both struct-field tests in `KeyGroupedPartitioningSuite` run
without `allowCompatibleTransforms` as well. So a regression that lets
`bucket(8, s.a)` reduce onto `bucket(4, t.b)` passes every test. How about the
same `reducers` and `isCompatible` assertions for a struct field, e.g. with
`bucket(DivisibleBucketFunction, GetStructField(s, 0), 8)`?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -117,7 +127,10 @@ case class TransformExpression(
*/
def reducers(other: TransformExpression): Option[Reducer[_, _]] = {
(function, other.function) match {
- case(e1: ReducibleFunction[_, _], e2: ReducibleFunction[_, _]) =>
+ // A reduce pairs the two sides up the same way `isSameFunction` does,
so it needs the same
+ // arguments.
+ case (e1: ReducibleFunction[_, _], e2: ReducibleFunction[_, _])
+ if argumentsAreAttributes && other.argumentsAreAttributes =>
Review Comment:
**Could the pairing rule live in `KeyedShuffleSpec.isExpressionCompatible`
instead?** The rule exists because `KeyedShuffleSpec.keyPositions` pairs the
two sides by reference. But it is written into the general
`TransformExpression` API twice, in `isSameFunction` and here, and
`isCompatible` is rewritten to inherit it. As a result, a transform of an
expression is not the same function as itself, although the name of
`isSameFunction` and the `TransformFunctionId` doc promise function identity. A
single guard in the `(TransformExpression, TransformExpression)` case of
`KeyedShuffleSpec.isExpressionCompatible` gives the same results on every
production path:
- In production code, `isSameFunction` and `isCompatible` are called only
from `isExpressionCompatible`.
- `reducersBothWays`, the only other caller of `reducers`, is called only
for a pair that `areKeysCompatible(allowReduce = true)` admitted, at
`EnsureRequirements` L705 and L764.
It would also sit next to the existing rule there that admits only leaf
expressions, and it is just as small to backport. I understand that a check
here also protects future callers, so please take this as a suggestion.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -92,19 +113,8 @@ case class TransformExpression(
* @param other the transform expression to compare to
* @return true if compatible, false if not
*/
- def isCompatible(other: TransformExpression): Boolean = {
- if (isSameFunction(other)) {
- true
- } else {
- (function, other.function) match {
- case (f: ReducibleFunction[_, _], o: ReducibleFunction[_, _]) =>
- val thisReducer = reducer(f, numBucketsOpt, o, other.numBucketsOpt)
- val otherReducer = reducer(o, other.numBucketsOpt, f, numBucketsOpt)
- thisReducer.isDefined || otherReducer.isDefined
- case _ => false
- }
- }
- }
+ def isCompatible(other: TransformExpression): Boolean =
+ isSameFunction(other) || reducers(other).isDefined ||
other.reducers(this).isDefined
Review Comment:
**The Scaladocs still describe the old contract.** The `isCompatible`
Scaladoc right above says it is true if both functions are `ReducibleFunction`s
and a `Reducer` exists, and the `reducers` Scaladoc says a reducer exists
whenever one function is reducible on the other. The new test contradicts both:
`bucket8.reducers(bucket4)` is defined, but
`bucket8OfExpression.isCompatible(bucket4)` is false for the same functions. In
addition, `ShuffleSpec.isCompatibleWith` still says "Spark assumes this to be
reflexive, symmetric and transitive". It already had an exception in
`RangeShuffleSpec`, and now a spec over `bucket(4, b + 1)` is incompatible with
itself too. `pickCoPartitionTarget` avoids such a winner only because of the
new `canCreatePartitioning` clause. Could you mention the argument condition in
these docs?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/TransformExpressionSuite.scala:
##########
@@ -41,6 +41,24 @@ class TransformExpressionSuite extends SparkFunSuite {
override def hashCode(): Int = canonicalName().hashCode
}
+ /** Reduces a bucket count onto any smaller count that divides it. */
+ private object DivisibleBucketFunction
Review Comment:
nit: `ShuffleSpecSuite.FakeBucket` already reduces a bucket count onto a
divisor with `% otherNumBuckets`, and its doc says it is local only because
catalyst has no `BucketFunction`. Moving it next to `FlipLowBitFunction` in
`sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/functions/`
would let both suites share it. Also, only the positive branch of this
`reducer` runs, since the new guard in `reducers` refuses the other cases
before the function is called.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -2199,8 +2203,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,
Review Comment:
**The underlying issue is that `createPartitioning` replaces the whole
argument (L2225).** This clause refuses a template because the rewrite drops
whatever surrounds the key. Substituting only the reference that `keyPositions`
mapped, as `withReference` and the `IdentityReducer` case already do, would be
sound for a single-reference argument. `bucket(4, b + 1)` would become
`bucket(4, y + 1)` for `xy`, and `bucket(4, s.a)` would become `bucket(4,
t.a)`, which also removes the `ClassCastException`. It is not a drop-in
replacement for this clause, though:
- To get the conceded shuffles back, `isSameFunction` would also have to
compare the shapes and stay reflexive, since `pickCoPartitionTarget` keeps the
winning child in place only when its member pairs with itself.
- The shuffled side would evaluate `y + 1` on rows where the query never
computes it, and that can overflow under ANSI mode.
So I'm raising it as input for SPARK-59900 rather than as a change to this
PR.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -1520,6 +1520,113 @@ 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, which
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`, so pairing the two 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`,
+ // so a row with `b = y` lands in a different bucket on each side.
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ val df = sql(
+ s"""SELECT k.id, p.b, q.x FROM testcat.ns.plain p
+ |JOIN testcat.ns.keys k ON $key = k.id
+ |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))))
+ }
+ }
+ }
+
+ test("SPARK-59887: a side shuffled onto a transform of a struct field is not
paired as it is") {
+ createBucketedIdTable("bucket4", 4, numIds = 8)
+ val structColumn = Array(
+ Column.create("s", new StructType().add("a", LongType).add("b",
LongType)))
+ Seq("ps", "qs").foreach { name =>
+ createTable(name, structColumn, Array.empty)
+ sql(s"INSERT INTO testcat.ns.$name VALUES " +
+ (0 until 8).map(i => s"(named_struct('a', ${i}L, 'b', ${7 -
i}L))").mkString(", "))
+ }
+
+ // Each side is shuffled onto the bucket of a different field of `s`,
`bucket(4, s.a)` and
+ // `bucket(4, s.b)`. Both 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.
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ val df = sql(
+ """SELECT l.id, r.id FROM
+ | (SELECT x.id, ps.s FROM testcat.ns.bucket4 x JOIN testcat.ns.ps
ps ON x.id = ps.s.a) l
+ | JOIN
+ | (SELECT y.id, qs.s FROM testcat.ns.bucket4 y JOIN testcat.ns.qs
qs ON y.id = qs.s.b) r
+ | ON l.s = r.s""".stripMargin)
+ checkAnswer(df, (0 until 8).map(i => Row(i.toLong, (7 - i).toLong)))
+ }
+ }
+
+ test("SPARK-59887: a side is not shuffled onto a transform of a struct
field") {
Review Comment:
nit: This test repeats the setup of the previous one (`bucket4` with 8 ids
and the same `named_struct` rows) and its `l` subquery. Per the ablation in the
PR description, it fails only when the `canCreatePartitioning` clause is
removed, and the previous test fails then too. If you want to keep this
`ClassCastException` shape, it could be a second query in the previous test,
`SELECT l.id FROM <the same l> JOIN testcat.ns.qs r ON l.s = r.s`, like
`createKeyTables()` shares the setup of the expression tests. That saves the
`rs` table and about 20 lines.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -65,6 +65,16 @@ case class TransformExpression(
lazy val functionId: TransformFunctionId =
TransformFunctionId(function.canonicalName(), numBucketsOpt)
+ /**
+ * Whether every argument is a bare column, an [[Attribute]]. See
[[isSameFunction]] for why only
+ * such a transform can be compared with another one. A literal argument is
a parameter of the
+ * transform rather than a column, so it does not count.
+ */
+ private[sql] lazy val argumentsAreAttributes: Boolean = children.forall {
+ case _: Literal => true
Review Comment:
**The `Literal` case is unreachable, and its rationale doesn't hold for
either consumer.** No transform with a literal child reaches `isSameFunction`,
`reducers` or `canCreatePartitioning`. A scan's transform has to pass
`KeyedPartitioning.supportsExpressions`, which requires a single reference
child, and `createPartitioning` keeps the arity and fills the argument with a
join key. If one did arrive, e.g. `truncate(a, 2)`, both consumers would get it
wrong. `functionId` ignores the literal's value, so `truncate(a, 2)` would be
the same function as `truncate(b, 5)`. And `createPartitioning` would overwrite
the literal and build `truncate(x, x)`, which is the hazard the new comment in
`canCreatePartitioning` describes. How about
`children.forall(_.isInstanceOf[Attribute])`, without the sentence about
literals? It is equivalent today, and it fails safe.
--
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]
