peter-toth commented on code in PR #58335:
URL: https://github.com/apache/spark/pull/58335#discussion_r3871285779
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -70,11 +70,22 @@ case class GroupPartitionsExec(
// There can be multiple `KeyedPartitioning`s in an output
partitioning of a join, but they
// can only differ in `expressions`; their `partitionKeys` reference
is shared (enforced by
// `PartitioningCollection`), so `groupedPartitions` is computed only
once.
+ // When reducers are applied, the reduced expressions (whose data type
matches the reduced
+ // partition keys) are reported instead of the original ones.
val partitionKeys = groupedPartitions.map(_._1)
p.transform {
case k: KeyedPartitioning =>
val projectedExpressions =
joinKeyPositions.fold(k.expressions)(_.map(k.expressions))
- KeyedPartitioning(projectedExpressions, partitionKeys, isGrouped =
isGrouped)
+ val effectiveExpressions = reducers match {
+ case Some(exprs) =>
+ assert(projectedExpressions.length == exprs.length)
+ projectedExpressions.zip(exprs).map {
+ case (expr, Some((_, reduced))) => reduced
Review Comment:
**Finding 2.** `reduced` was derived from one `KeyedShuffleSpec`, but this
`transform` applies it to every `KeyedPartitioning` in the child's
partitioning, so the others lose their own key attribute.
`createKeyedShuffleSpec` takes the first `KeyedPartitioning` that satisfies
the distribution (`collectFirst`, `EnsureRequirements.scala:819`), and the
reducers are computed from that one alone. When this node sits on a join - a
chained SPJ - the child reports one `KeyedPartitioning` per side, and all of
them get that single `reduced` expression, which references only the chosen
side's attribute. The `assert` above has the same root: the length is only
guaranteed to match for the spec's own `KeyedPartitioning`.
Three tables bucketed 16, 8 and 4 on `id`, ids 0..15, joined on `a.id` and
grouped by `b.id`:
- base: `KPs = [bucket(16, id#18L), bucket(8, id#20L)]`, 0 shuffles.
- this commit: `KPs = [bucket(16, id#18L), bucket(16, id#18L)]`, 1 shuffle,
because `GROUP BY b.id` no longer sees a partitioning on `b.id`.
Both return the same rows, so this is an avoidable shuffle rather than a
wrong answer. Retargeting the reduced expression at each `KeyedPartitioning`'s
own key attribute fixes it, and the single-attribute invariant makes that well
defined:
```scala
projectedExpressions.zip(exprs).map {
case (expr, Some((_, reduced))) =>
// `reduced` was derived from the spec's `KeyedPartitioning`; re-target
it at this one's key so
// that every `KeyedPartitioning` in a collection keeps its own
attribute.
val attr = expr.references.head
reduced.transform { case _: AttributeReference => attr }
case (expr, None) => expr
}
```
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1374,9 +1376,11 @@ case class KeyedShuffleSpec(
*
* @param other other key-grouped shuffle spec
*/
- def reducers(other: KeyedShuffleSpec): Option[Seq[Option[Reducer[_, _]]]] = {
+ def reducers(
+ other: KeyedShuffleSpec): Option[Seq[Option[(Reducer[_, _],
TransformExpression)]]] = {
val results =
partitioning.expressions.zip(other.partitioning.expressions).map {
- case (e1: TransformExpression, e2: TransformExpression) =>
e1.reducers(e2)
+ case (e1: TransformExpression, e2: TransformExpression) =>
+ e1.reducers(e2).map(reducer => (reducer, e1))
Review Comment:
**Finding 3.** This branch covers two different shapes, and keeping `e1` is
wrong in both. Not asking you to fix them here - the second one needs a design
decision - but they should not be left implicit either.
There are three reducer shapes, and the reported expression has to be judged
per shape:
1. **This side is identity, the other is a transform.** The reduced
expression exists and the branch below builds it. Handled, and correctly.
2. **Both sides are transforms, only this side reduces.** The reduced
expression also exists: `r(f1(x)) = f2(x)` is the `ReducibleFunction` contract,
so it is the other side's transform retargeted at this side's child. This
branch reports `e1` instead. `SPARK-56046: Reducers with same result types` is
this shape - `days` reduces onto `years`, `YearsFunction.reducer` returns null
- and the reported expression stays `days(arrive_time)` where
`years(arrive_time)` is what the keys hold.
3. **Both sides reduce.** The keys are `r1(f1(x)) = r2(f2(x))`, a space that
neither transform describes, so no substitution can be correct here.
Shape 3 still throws the exception this PR fixes. `reduceKeys` types the
keys with `reducer.resultType()` while `e1.dataType` is
`e1.function.resultType()`, and the suite already ships a pair where they
differ: `DaysFunctionWithToYearsReducerWithLongResult` (`DateType`) and
`YearsFunctionWithToYearsReducerWithLongResult` (`IntegerType`) both reduce to
`LongType`. Take `SPARK-56164: Reducers with different result types to original
keys` and add `V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS -> true`:
```
java.lang.ClassCastException: class java.lang.Long cannot be cast to class
java.lang.Integer
at
org.apache.spark.sql.catalyst.expressions.GenericInternalRow.getInt(rows.scala:170)
at
org.apache.spark.sql.catalyst.expressions.GeneratedClass$SpecificOrdering.compare(Unknown
Source)
at
org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning.toGrouped(partitioning.scala:553)
at
org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning.createShuffleSpec(partitioning.scala:629)
at
org.apache.spark.sql.execution.exchange.ValidateRequirements$.validateInternal(ValidateRequirements.scala:60)
```
`e2` would not help: `years` is `IntegerType` and the keys are `LongType`.
The other cost of a stale expression is that a *second* join derives its
reducer from a transform the data has already left behind. Three tables
bucketed 12, 8 and 6 on `id`, ids 0..11: join 1 reduces both sides with
`BucketReducer(4)`, so the keys become `id % 4` while the left side reports
`bucket(12, id)`; join 2 compares that stale `bucket(12, id)` with `bucket(6,
id)`, gets `BucketReducer(6)`, and `(id % 4) % 6` leaves the keys at `id % 4`,
while the right side has `gcd == thisNumBuckets`, gets no reducer, and keeps
`id % 6`. The two sides are matched across different key spaces: the query
returns 4 of 12 rows, 0 shuffles, no error, and all 12 with
`allowCompatibleTransforms=false`.
```scala
val cols = Array(Column.create("id", LongType), Column.create("data",
StringType))
createTable("b12", cols, Array(bucket(12, "id")))
createTable("b8", cols, Array(bucket(8, "id")))
createTable("b6", cols, Array(bucket(6, "id")))
val values = (0 until 12).map(i => s"($i, 'v$i')").mkString(", ")
Seq("b12", "b8", "b6").foreach(t => sql(s"INSERT INTO testcat.ns.$t VALUES
$values"))
val df = sql(
"""SELECT /*+ MERGE(a, b, c) */ a.id FROM testcat.ns.b12 a
|JOIN testcat.ns.b8 b ON a.id = b.id
|JOIN testcat.ns.b6 c ON a.id = c.id""".stripMargin)
withSQLConf(SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
checkAnswer(df, (0 until 12).map(i => Row(i.toLong))) // returns only 0,
1, 2, 3
}
```
That one is pre-existing, base reports the same expression, so it is not
something this PR broke. It is not master-only either: the same query returns 8
of 12 rows on `branch-4.1`, where
`KeyGroupedPartitionedScan.getOutputKeyGroupedPartitioning` reports the
original expressions next to the reduced common partition values in the same
way.
The shape I would suggest, in case it is useful: extend the substitution to
shape 2, which is a small step from what you already do for shape 1; and for
shape 3 stop trying to express the keys as a transform - carry the reduced key
data types on `KeyedPartitioning` instead. `reduceKeys` already computes them
one frame up and `GroupPartitionsExec.groupedPartitionsTuple` already holds
them as `reducedDataTypes`, they are just dropped, and the reported
partitioning re-derives its types from `expressions.map(_.dataType)`
(`partitioning.scala:545`). All eight readers go through `expressionDataTypes`,
so they would pick the carried types up for free. Such a partitioning then has
to refuse a further reduction, since its expressions no longer describe its
keys - which is exactly what the 12/8/6 case above needs.
Happy to pick this up as a follow-up if you would rather not carry it - tell
me which way you prefer.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -542,6 +542,87 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase with
}
}
+ test("SPARK-59045: compatible identity and bucket transforms reduce data
type") {
+ // `identity(id)` reports a Long partition key while `bucket(4, id)`
reports an Integer one.
+ // The identity->bucket reducer maps the Long keys to Integer; the
GroupPartitionsExec output
+ // partitioning must report the reduced (Integer) expression, not the
original Long identity,
+ // or the key ordering derived from the expressions fails with a
ClassCastException.
+ val cols = Array(
+ Column.create("id", LongType),
+ Column.create("data", StringType))
+ createTable("t1", cols, Array(identity("id")))
+ sql("INSERT INTO testcat.ns.t1 VALUES (1, 'a'), (2, 'b'), (3, 'c')")
+
+ createTable("t2", cols, Array(bucket(4, "id")))
+ sql("INSERT INTO testcat.ns.t2 VALUES (1, 'x'), (2, 'y'), (3, 'z')")
+
+ val df = sql(
+ "SELECT t1.id, t1.data, t2.data FROM testcat.ns.t1 JOIN testcat.ns.t2 ON
t1.id = t2.id")
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true") {
+ checkAnswer(df, Seq(Row(1, "a", "x"), Row(2, "b", "y"), Row(3, "c",
"z")))
+ assert(collectShuffles(df.queryExecution.executedPlan).isEmpty,
+ "storage-partitioned join should not shuffle")
+ }
+ }
+
+ test("SPARK-59045: compatible transforms reduce multiple times") {
+ // t1 is partitioned by identity(id) (Long), t2 by bucket(4, id), t3 by
bucket(2, id). The
+ // first join reduces t1 to bucket(4, id) (data type changes), and the
second join reduces the
+ // result to bucket(2, id). The reduced expression reported by the first
join must remain a
+ // ReducibleFunction so the second reduction can be computed.
+ val cols = Array(Column.create("id", LongType), Column.create("data",
StringType))
+ createTable("t1", cols, Array(identity("id")))
+ createTable("t2", cols, Array(bucket(4, "id")))
+ createTable("t3", cols, Array(bucket(2, "id")))
+ sql("INSERT INTO testcat.ns.t1 VALUES (1, 'a'), (2, 'b'), (3, 'c')")
+ sql("INSERT INTO testcat.ns.t2 VALUES (1, 'x'), (2, 'y'), (3, 'z')")
+ sql("INSERT INTO testcat.ns.t3 VALUES (1, 'p'), (2, 'q'), (3, 'r')")
+
+ val df = sql(
+ "SELECT t1.id, t1.data, t2.data, t3.data FROM testcat.ns.t1 " +
+ "JOIN testcat.ns.t2 ON t1.id = t2.id JOIN testcat.ns.t3 ON t1.id =
t3.id")
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key ->
"true") {
Review Comment:
**Finding 4.** All three new tests enable this config, so all three fail on
base at the same place, `createShuffleSpec` -> `toGrouped`. There is a second,
independent trigger that none of them reach.
This test does not need the config, the join uses the whole partition key.
With it removed the test still fails on base, but through `reduceKeys` at the
second join, where the reducer is bound to the stale un-reduced type:
```
java.lang.ClassCastException: class java.lang.Integer cannot be cast to
class java.lang.Long
at
org.apache.spark.sql.catalyst.plans.physical.KeyedShuffleSpec$$anon$1.reduce(partitioning.scala:1394)
at
org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning$.reduceKeys(partitioning.scala:725)
at
org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning.reduceKeys(partitioning.scala:572)
at
org.apache.spark.sql.execution.exchange.EnsureRequirements.checkKeyGroupCompatible(EnsureRequirements.scala:561)
```
and it passes on this commit. Test 1 and test 3 do need the config, test 3
genuinely joins on a subset.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1374,9 +1376,11 @@ case class KeyedShuffleSpec(
*
* @param other other key-grouped shuffle spec
*/
- def reducers(other: KeyedShuffleSpec): Option[Seq[Option[Reducer[_, _]]]] = {
+ def reducers(
+ other: KeyedShuffleSpec): Option[Seq[Option[(Reducer[_, _],
TransformExpression)]]] = {
Review Comment:
**Finding 6.** `Option[Seq[Option[(Reducer[_, _], TransformExpression)]]]`
now appears in six signatures and forces `_1` at the use sites, including the
explain string. A small named type would read better and would document which
element is which:
```scala
case class KeyReducer(reducer: Reducer[_, _], reducedExpression:
TransformExpression)
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -70,11 +70,22 @@ case class GroupPartitionsExec(
// There can be multiple `KeyedPartitioning`s in an output
partitioning of a join, but they
// can only differ in `expressions`; their `partitionKeys` reference
is shared (enforced by
// `PartitioningCollection`), so `groupedPartitions` is computed only
once.
+ // When reducers are applied, the reduced expressions (whose data type
matches the reduced
Review Comment:
**Finding 5.** The parenthetical holds for the identity branch only. For two
transforms `reducers` carries `e1`, so the reported data type is the un-reduced
one and can differ from the keys (finding 3). Worth naming the branch the claim
applies to, since the next reader will lean on 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]
