metanil commented on code in PR #59760:
URL: https://github.com/apache/spark/pull/59760#discussion_r4432840456
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -1792,6 +1792,77 @@ class KeyGroupedPartitioningSuite
}
}
+ test("SPARK-60522: a transform Spark cannot evaluate is not reduced onto or
shuffled onto") {
+ // With `UnboundBucketFunction`, the first join reduces `ident` onto
`bucket(4, id)`, and the
+ // second shuffles `plain` onto it, both by evaluating `bucket`. A
connector may bind `bucket`
+ // to a function that is not a `ScalarFunction`, or to one Spark finds
nothing to call on, and
+ // Spark cannot evaluate either. Then both joins shuffle both sides.
+ Seq(
+ UnboundBucketFunction -> (0, 1),
+ UnboundNonScalarBucketFunction -> (2, 2),
+ UnboundInheritedBucketFunction -> (2, 2)
+ ).foreach {
+ case (fn, (reduceShuffles, oneSideShuffles)) =>
+ withFunction(fn) {
+ createBucketedIdTable("bucket4", 4, numIds = 8)
+ createIdentityIdTable()
+ createPlainTable()
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ val df = sql(
+ "SELECT k.id FROM testcat.ns.bucket4 k JOIN testcat.ns.ident i
ON k.id = i.id")
+ checkAnswer(df, (0 until 8).map(i => Row(i.toLong)))
+ assert(collectAllShuffles(df.queryExecution.executedPlan).size ==
reduceShuffles,
+ fn.name())
+ }
+
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ val df = sql(
+ "SELECT k.id FROM testcat.ns.bucket4 k JOIN testcat.ns.plain p
ON k.id = p.b")
+ checkAnswer(df, (0 until 8).map(i => Row(i.toLong)))
+ assert(collectAllShuffles(df.queryExecution.executedPlan).size ==
oneSideShuffles,
+ fn.name())
+ }
+ catalog.clearTables()
+ }
+ }
+ }
+
+ test("SPARK-60522: a child Spark cannot evaluate leaves the other child as
the shuffle target") {
+ // `bucket8` lives in a second catalog, whose `bucket` Spark can evaluate.
Picking the target,
+ // `EnsureRequirements` asks every child whether it can serve as one, and
`bucket4`'s must
+ // answer no, not throw, so that `bucket4` is shuffled onto `bucket8`.
+ val catalogName = "spark60522cat"
+ spark.conf.set(s"spark.sql.catalog.$catalogName",
classOf[InMemoryCatalog].getName)
+ val catalog2 = spark.sessionState.catalogManager.catalog(catalogName)
+ .asInstanceOf[InMemoryCatalog]
+ try {
+ withFunction(UnboundInheritedBucketFunction) {
+ catalog2.createFunction(Identifier.of(Array.empty, "bucket"),
UnboundBucketFunction)
+ createBucketedIdTable("bucket4", 4, numIds = 8)
+ createTable("bucket8", Array(Column.create("id", LongType)),
Array(bucket(8, "id")),
+ catalog = catalog2)
+ sql(s"INSERT INTO $catalogName.ns.bucket8 VALUES " +
+ (0 until 8).map(i => s"($i)").mkString(", "))
+
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ Seq(s"testcat.ns.bucket4 a JOIN $catalogName.ns.bucket8 b",
Review Comment:
Good catch, thanks. The test now also runs with the counts swapped: the
inherited-`produceResult` `bucket(8) `against a scalar `bucket(4)` in the
second catalog. That's the case that fails on the base, so the test now
actually exercises the gate.
--
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]
