ulysses-you commented on code in PR #58335:
URL: https://github.com/apache/spark/pull/58335#discussion_r3871588965
##########
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:
Fixed in 2b0ef181bc4. `GroupPartitionsExec.outputPartitioning` now
re-targets the reduced expression at each `KeyedPartitioning`'s own key
attribute before applying it, so a chained SPJ keeps every side's partitioning
on its own attribute. Added the 16/8/4 + GROUP BY middle-side regression test;
it asserts 0 shuffles and the full 16-row result.
##########
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:
Thanks for the detailed breakdown. Shape 2 is now handled in 2b0ef181bc4:
for a single-side reduce we report the target transform re-targeted at this
side's attribute, guarded by `e2.reducers(e1).isEmpty`, so both-sides-reduce
positions keep the original expression bit-for-bit (shape 3 unchanged, still a
tracked gap).
##########
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:
Done in 2b0ef181bc4. Dropped
`V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS` from the multi-table test
(both joins use the whole partition key); it now reaches the `reduceKeys`
trigger and fails on base at `KeyedShuffleSpec.reduce`. Added a comment noting
the two triggers.
##########
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:
Fixed in 2b0ef181bc4. The comment now names the branches: the data-type
match holds for the identity-vs-transform and single-side-transform reducers,
and the both-sides-reduce shape has no single transform expression (pointed at
`KeyedShuffleSpec.reducers`).
##########
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:
Done in 2b0ef181bc4. Added `case class KeyReducer(reducer,
reducedExpression)` next to `KeyedShuffleSpec` and switched all signatures/use
sites (including the explain string) to 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]
