peter-toth commented on code in PR #59760:
URL: https://github.com/apache/spark/pull/59760#discussion_r4370153120
##########
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:
**Finding 3.** `pickCoPartitionTarget` ranks the candidates by partition
count, so `bucket8` is the target here whether or not `bucket4` could serve as
one. So the test passes on the base, as the description says. It also passes
with `canEvaluate` taken out of `canCreatePartitioning`. It pins only that the
gate does not throw.
I swapped the counts: `bucket8` in `testcat` over
`UnboundInheritedBucketFunction`, and `bucket4` in the second catalog over
`UnboundBucketFunction`. Both join orders then return 8 rows with one shuffle
here. On the base they fail in a task with
`SCALAR_FUNCTION_NOT_FULLY_IMPLEMENTED`. Adding that pair would make the test
pin the choice its name describes.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -2570,8 +2592,8 @@ case class KeyedShuffleSpec(
// `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 => argumentsAreColumns(t)
+ // children shuffled onto it. Those children are routed by evaluating
it, too.
+ case t: TransformExpression => argumentsAreColumns(t) && canEvaluate(t)
Review Comment:
**Finding 1.** `pickCoPartitionTarget` (`EnsureRequirements.scala:376-423`)
asks `canCreatePartitioning` for two things:
- which layout the other children are shuffled onto,
- which layout a child is read against as it stands.
Only the first one evaluates the transform. An operator that co-partitions
but is not a sort-merge or shuffled hash join has only this path, since
`checkKeyGroupCompatible` declines for it. That is a Pandas or Arrow cogroup,
or `SortMergeAsOfJoinExec`.
With `spark.sql.sources.v2.bucketing.shuffle.enabled`, take two sides that
report the same non-evaluable `bucket(4, id)` and the same keys. The base reads
them as they stand, and nothing evaluates the transform. With this PR neither
side can be the target, so both get a hash shuffle. I measured it in
`EnsureRequirementsSuite`, on the shape of the positive control in
"SPARK-59688: a cogroup does not pair two key spaces as they stand", with
`nonScalar` a `BoundFunction` that is not a `ScalarFunction`:
```scala
val keys = Seq(InternalRow(0), InternalRow(1), InternalRow(2),
InternalRow(3))
val left = DummySparkPlan(outputPartitioning = KeyedPartitioning(
Seq(TransformExpression(nonScalar, Seq(idL), Some(4))), keys))
val right = DummySparkPlan(outputPartitioning = KeyedPartitioning(
Seq(TransformExpression(nonScalar, Seq(idR), Some(4))), keys))
val cogroup = FlatMapCoGroupsInPandasExec(Seq(idL), Seq(idR), pythonUdf,
AttributeReference("value", IntegerType)() :: Nil, left, right)
withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
// ShuffleExchangeExec count: 0 on the base, 2 with this PR
EnsureRequirements.apply(cogroup)
}
```
A `SortMergeJoinExec` over the same two children stays at 0 shuffles on
both, since `checkKeyGroupCompatible` pairs it first.
The rows are right either way, so this is a lost pairing, not a wrong
result. That pairing is already tied to `canCreatePartitioning`, so it is lost
with `shuffle.enabled` off too. Keeping it would mean telling "can be read as
it stands" apart from "can be shuffled onto" in `pickCoPartitionTarget`. I
think that can wait for a follow-up. Could you list this case under the
user-facing changes, though? The snippet above can pin it as a test.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -2415,6 +2419,23 @@ case class KeyedShuffleSpec(
private def argumentsAreColumns(transform: TransformExpression): Boolean =
transform.children.forall(_.isInstanceOf[Attribute])
+ /**
+ * Whether Spark can evaluate `transform` itself. An identity reduce
evaluates it on the identity
+ * side's partition keys, and a one-side shuffle on each row of the side it
shuffles. Only a
+ * `ScalarFunction` can be evaluated (see
`TransformExpression.resolvedFunction`), while a
+ * connector may bind its transform to any `BoundFunction`. Rebuilding the
transform over
+ * another key keeps its function, so asking this of the reported transform
answers it for the
+ * rebuilt one too.
+ *
+ * Resolving a `ScalarFunction` throws when Spark finds no method to call on
it, e.g. one that
+ * inherits `produceResult` from a base class. That counts as cannot
evaluate, not as a failure:
+ * the gates ask this of every candidate, including ones that end up not
being chosen.
+ */
+ private def canEvaluate(transform: TransformExpression): Boolean =
+ try transform.resolvedFunction.isDefined catch {
Review Comment:
**Finding 2.** `resolveScalarFunction` looks `produceResult` up with
`getDeclaredMethod`. A Scala class or object that extends `ScalarFunction`
directly gets a mixin forwarder to the Java default `produceResult`. So the
lookup succeeds even when the connector implements neither `produceResult` nor
`invoke`, and `canEvaluate` returns true. Evaluating it then hits the default,
which throws `SCALAR_FUNCTION_NOT_COMPATIBLE`. The TODO at
`V2ExpressionUtils.scala:231` describes the same lookup gap.
The gap predates this PR, and reflection can't tell the forwarder from a
real override, so I wouldn't try to fix it here. Could you name this case in
the `canEvaluate` doc, and in the `reducersBothWays` comment at L2547, which
says Spark can evaluate `t`?
The `"scalar"` control in the new `ShuffleSpecSuite` test is such a
function. `TestBucketFunction` implements neither method, and `javap` shows the
forwarder on `ShuffleSpecSuite$TestBucketFunction$`. Its row asserts `true`
only through the forwarder, and it would flip if that lookup gap were closed.
Giving `TestBucketFunction` a `produceResult` makes the control a function
Spark can run.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -2570,8 +2592,8 @@ case class KeyedShuffleSpec(
// `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 => argumentsAreColumns(t)
+ // children shuffled onto it. Those children are routed by evaluating
it, too.
+ case t: TransformExpression => argumentsAreColumns(t) && canEvaluate(t)
Review Comment:
**Finding 4.** For a partitioning whose keys a join reduced, this reads
`resolvedFunction` before `expressionsDescribeKeys` below refuses it. That
reaches the `reducedWith` arm, and the doc at
`TransformExpression.scala:173-175` says no caller reaches it. Asking
`expressionsDescribeKeys` first would keep that doc true. Otherwise the doc
needs updating.
--
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]
