peter-toth commented on code in PR #59165:
URL: https://github.com/apache/spark/pull/59165#discussion_r4152407781
##########
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:
Thanks, confirmed, and fixed in
[69f16ab](https://github.com/apache/spark/pull/59165/commits/69f16abd834a271f68e04905d5391b887fcb4a49)
the way you suggest. `GroupPartitionsExec` now rebuilds `reduced` over the
member's own argument, the child of a transform or the member itself. So a
`bucket(4, b + 1)` member becomes `bucket(2, b + 1)`, which the guard refuses.
`reduced`'s own argument is always a bare column now, since the guard also
covers the identity arm
([thread](https://github.com/apache/spark/pull/59165#discussion_r4148432898)).
"SPARK-59887: a reduce keeps what each member's keys are computed from" runs
your four shapes. It returns 0 of 7 rows without the change.
##########
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:
Fixed in
[69f16ab](https://github.com/apache/spark/pull/59165/commits/69f16abd834a271f68e04905d5391b887fcb4a49),
as SPARK-59901, which I had filed for this and have now added to the title.
`keyPositions` maps an expression with more than one reference to no position
instead of asserting. `canCreatePartitioning` and `areKeysCompatible` both turn
an empty position away, so a spec over `b + c` pairs with nothing and is no
layout to shuffle onto. Both of your queries are in "SPARK-59901: a join key
over two columns is not paired through one of them".
##########
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:
Confirmed. It is a separate consumer, so I filed SPARK-59905 for it rather
than growing this PR. The same binding also fails planning for `ORDER BY p.b +
p.c` over the two-column key, with an `ArrayIndexOutOfBoundsException`, so I
added that shape to the ticket. The `EnsureRequirements` comment that claimed
one attribute per partition expression for this binding now points there.
##########
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:
Thanks, done in
[69f16ab](https://github.com/apache/spark/pull/59165/commits/69f16abd834a271f68e04905d5391b887fcb4a49).
`ShuffleSpecCollection.canCreatePartitioning` is now an `exists`, and
`pickCoPartitionTarget` offers only the winner's members that can serve.
`PartitioningCollection.createShuffleSpec` admits only members that could
satisfy the distribution, so a collection without a keyed member hardly ever
mixes the two answers. The planner, join, bucketed read, AQE and plan stability
suites pass unchanged. Your two shapes and the `xy` one take 2 shuffles again,
and the description no longer lists them as costs.
##########
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:
Right. I added it to the costs in the description. I'd rather not treat an
up-cast as the column in a fix meant for the maintenance branches. That leans
on each connector's canonical name meaning the same across input types.
##########
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:
Good point, done in
[69f16ab](https://github.com/apache/spark/pull/59165/commits/69f16abd834a271f68e04905d5391b887fcb4a49).
The test now runs both join orders, and it asserts 2 shuffles with none keyed
on `y`.
##########
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:
Added in
[69f16ab](https://github.com/apache/spark/pull/59165/commits/69f16abd834a271f68e04905d5391b887fcb4a49).
The unit test moved to `ShuffleSpecSuite` together with the guard, and it now
runs every assertion, the reduce in both directions included, for the struct
field as well.
##########
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:
Agreed, moved in
[69f16ab](https://github.com/apache/spark/pull/59165/commits/69f16abd834a271f68e04905d5391b887fcb4a49).
`TransformExpression` is back to master. The guard is
`KeyedShuffleSpec.argumentsAreColumns`, asked at the entry of
`isExpressionCompatible` and by `canCreatePartitioning`.
At the entry it also covers the identity arm. That arm used to reduce an
identity side onto `bucket(4, b + 1)` by evaluating `bucket(4, id + 1)` on the
identity side's partition keys while planning. Under ANSI, a key of
`Long.MaxValue` then fails the query with `ARITHMETIC_OVERFLOW`, although the
query never computes `id + 1`. "SPARK-59887: an identity side is not reduced
onto a transform of an expression" covers it.
##########
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:
Thanks, this is very useful for SPARK-59900, the ANSI point especially.
Substituting the reference makes the shuffled side evaluate `y + 1` on rows
where the query never computes it.
--
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]
