peter-toth commented on code in PR #59165:
URL: https://github.com/apache/spark/pull/59165#discussion_r4152408936
##########
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:
Done in
[69f16ab](https://github.com/apache/spark/pull/59165/commits/69f16abd834a271f68e04905d5391b887fcb4a49).
With the guard out of `TransformExpression`, its docs hold again as they were.
The `ShuffleSpec.isCompatibleWith` note now lists the specs that are not
reflexive, and says what keeps the planner from picking one as a layout.
##########
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:
Done in
[69f16ab](https://github.com/apache/spark/pull/59165/commits/69f16abd834a271f68e04905d5391b887fcb4a49).
The check is `children.forall(_.isInstanceOf[Attribute])`, without the literal
case.
##########
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:
Merged into one test in
[69f16ab](https://github.com/apache/spark/pull/59165/commits/69f16abd834a271f68e04905d5391b887fcb4a49).
The `ClassCastException` shape is its second query, over a plain `ps`.
##########
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:
Gone in
[69f16ab](https://github.com/apache/spark/pull/59165/commits/69f16abd834a271f68e04905d5391b887fcb4a49).
The unit test moved to `ShuffleSpecSuite`, where it uses `FakeBucket`.
--
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]
