You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ul...@apache.org on 2026/09/02 01:27:38 UTC
(spark) branch branch-4.2 updated: [SPARK-59045][4.2][SQL] Fix SPJ ClassCastException when reducer changes partition key data type
This is an automated email from the ASF dual-hosted git repository.
ulysses-you pushed a commit to branch branch-4.2
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.2 by this push:
new 773f49a04568 [SPARK-59045][4.2][SQL] Fix SPJ ClassCastException when reducer changes partition key data type
773f49a04568 is described below
commit 773f49a0456875c2887c18043e8abd1adaff3de1
Author: Xiduo You <ul...@gmail.com>
AuthorDate: Wed Sep 2 09:27:22 2026 +0800
[SPARK-59045][4.2][SQL] Fix SPJ ClassCastException when reducer changes partition key data type
### What changes were proposed in this pull request?
Backport of #58335 to `branch-4.2`.
`KeyedShuffleSpec.reducersBothWays` now pairs each `Reducer` with the reduced partition expression it produces (`KeyReducer`), and `GroupPartitionsExec.outputPartitioning` reports the reduced expression instead of the original partition expressions when reducers are applied. The stored expression is re-targeted at each `KeyedPartitioning`'s own key attribute at the use site via the new `TransformExpression.withReference`, so a chained storage-partitioned join keeps every side's partit [...]
`GroupPartitionsExec.doCanonicalize` additionally normalizes the exprIds inside `KeyReducer` - plan canonicalization does not reach into the plain case class - and the reducer applied for an identity-vs-transform pair is a named `IdentityReducer` case class, so structurally identical SPJ subtrees with value-equal reducers still compare equal and exchange/subquery reuse keeps deduplicating them.
Tailored for this branch in four places:
- `GroupPartitionsExec.outputPartitioning` keeps this branch's `groupedPartitions`/`isGrouped` reporting and its multi-`KeyedPartitioning` partition-keys assertion. The `PartitionGrouping`/`isCollapsed` refactor, which replaced both on `master`, is not on this branch.
- The config is spelled `V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS` here, the `JOIN_` was dropped from the name later.
- The canonicalization test builds `LocalTableScanExec` with this branch's three-argument constructor.
- The subset-join-key test additionally sets `REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION` to false. SPARK-58558, which relaxed `createKeyedShuffleSpec` from an exact key match to a key-coverage check so the subset shape plans under the default, is not on this branch; the branch's other subset tests set the config the same way.
One of the three `SPARK-59120` tests from #58335 is not carried over: `SPARK-59120: reduced partition keys are read at the types they were built with` was already omitted from the #58420 backport (#58431), because its failure mode runs through the sort that `createShuffleSpec` applies, which arrived with SPARK-59022 and is not on this branch. The other two `SPARK-59120` tests are carried over verbatim; like on `master`, they replace the two #58431 tests and pin the combined behavior.
### Why are the changes needed?
In a storage-partitioned join with compatible transforms whose result types differ (e.g. `identity(id)` on one side and `bucket(N, id)` on the other), the reducer maps the partition keys to the other side's value type. `GroupPartitionsExec.outputPartitioning` used to report the original expressions with the reduced keys, so the two had different data types and computing the key ordering threw:
```
java.lang.ClassCastException: class java.lang.Integer cannot be cast to class java.lang.Long
```
This fix covers the reducers where the reduced keys equal a single transform applied to one side (identity-vs-transform and single-side-transform). When both sides of a compatible-transform join reduce their keys, the reduced keys are not expressible as a single transform; that shape is a known gap tracked in SPARK-59121.
### Does this PR introduce _any_ user-facing change?
No by default. Under `spark.sql.sources.v2.bucketing.allowCompatibleTransforms.enabled`, a storage-partitioned join whose reducer changes the partition key data type previously threw `ClassCastException` and now succeeds.
### How was this patch tested?
Added regression tests in `KeyGroupedPartitioningSuite`, mirroring #58335. Measured failing on this branch at `0e37b68e742` (with #58431 on it):
- `SPARK-59045: compatible transforms reduce multiple times` fails with `ClassCastException: class java.lang.Integer cannot be cast to class java.lang.Long` - the second reduce's reducer is derived from the stale reported expression, which no type fix reaches.
- The two carried-over `SPARK-59120` tests fail: the second-join one raises `STORAGE_PARTITION_JOIN_INCOMPATIBLE_REDUCED_TYPES` (the behavior #58431 pinned on this branch), and the another-child one plans two shuffles instead of one. They now pin the combined behavior.
- The `shipments` extension of `SPARK-56046: Reducers with same result types` fails when the third table reduces onto the first join's stale reported expression and that reduce evaluates the expression (`SCALAR_FUNCTION_NOT_FULLY_IMPLEMENTED`).
Measured passing on that base, as on `master`: the identity-vs-bucket reducer, the subset-join-key, and the per-`KeyedPartitioning` retargeting tests - they pin this change's reporting, and the retargeting test regresses together with the multi-reduce test if the use-site re-targeting is dropped. The canonicalization test arrived with `KeyReducer` in this change. The first three also assert the storage-partitioned join introduces no shuffle. All tests pass here.
Ran `KeyGroupedPartitioningSuite`, `GroupPartitionsExecSuite`, `EnsureRequirementsSuite`, `ProjectedOrderingAndPartitioningSuite`, `ShuffleSpecSuite`, `ValidateRequirementsSuite`, and `DistributionSuite`.
### Was this patch authored or co-authored using generative AI tooling?
Yes. Generated-by: Claude Code.
Closes #58451 from ulysses-you/spj-reducer-4.2.
Authored-by: Xiduo You <ul...@gmail.com>
Signed-off-by: Xiduo You <ul...@apache.org>
---
.../catalyst/expressions/TransformExpression.scala | 10 +
.../sql/catalyst/plans/physical/partitioning.scala | 142 ++++++++++---
.../spark/sql/errors/QueryExecutionErrors.scala | 10 +-
.../datasources/v2/GroupPartitionsExec.scala | 59 +++++-
.../execution/exchange/EnsureRequirements.scala | 6 +-
.../connector/KeyGroupedPartitioningSuite.scala | 225 ++++++++++++++++++---
6 files changed, 378 insertions(+), 74 deletions(-)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala
index 9041ed15fc50..437d394cb4a3 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala
@@ -98,6 +98,16 @@ case class TransformExpression(
}
}
+ /**
+ * Re-targets this partition transform expression at `attr`. A partition transform expression
+ * has a single leaf attribute (`KeyedPartitioning.supportsExpressions`), so this replaces that
+ * attribute and keeps the rest of the expression: any field path above the leaf comes from this
+ * expression, not from the re-targeted key. Re-targeting is therefore only faithful when the
+ * source and the target key expressions have the same path shape.
+ */
+ def withReference(attr: Attribute): TransformExpression =
+ transform { case _: AttributeReference => attr }.asInstanceOf[TransformExpression]
+
// Return a Reducer for a reducible function on another reducible function
private def reducer(
thisFunction: ReducibleFunction[_, _],
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
index b53648bfd611..e1482b4ae4f8 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
@@ -503,7 +503,7 @@ case class KeyedPartitioning(
* Returns the reduced keys and their data types.
*/
def reduceKeys(
- reducers: Seq[Option[Reducer[_, _]]]): (Seq[DataType], Seq[InternalRowComparableWrapper]) =
+ reducers: Seq[Option[KeyReducer]]): (Seq[DataType], Seq[InternalRowComparableWrapper]) =
KeyedPartitioning.reduceKeys(partitionKeys, keyDataTypes, reducers)
override def satisfies0(required: Distribution): Boolean = {
@@ -622,9 +622,9 @@ object KeyedPartitioning {
def reduceKeys(
keys: Seq[InternalRowComparableWrapper],
dataTypes: Seq[DataType],
- reducers: Seq[Option[Reducer[_, _]]]): (Seq[DataType], Seq[InternalRowComparableWrapper]) = {
+ reducers: Seq[Option[KeyReducer]]): (Seq[DataType], Seq[InternalRowComparableWrapper]) = {
val reducedDataTypes = dataTypes.zip(reducers).map {
- case (_, Some(reducer: Reducer[Any, Any])) => reducer.resultType()
+ case (_, Some(KeyReducer(reducer: Reducer[Any, Any], _))) => reducer.resultType()
case (t, _) => t
}
val comparableKeyWrapperFactory =
@@ -632,7 +632,7 @@ object KeyedPartitioning {
val reducedKeys = keys.map { key =>
val keyValues = key.row.toSeq(dataTypes)
val reducedKey = keyValues.zip(reducers).map {
- case (v, Some(reducer: Reducer[Any, Any])) => reducer.reduce(v)
+ case (v, Some(KeyReducer(reducer: Reducer[Any, Any], _))) => reducer.reduce(v)
case (v, _) => v
}.toArray
comparableKeyWrapperFactory(new GenericInternalRow(reducedKey))
@@ -975,6 +975,48 @@ case class CoalescedHashShuffleSpec(
override def numPartitions: Int = partitions.length
}
+/**
+ * A [[Reducer]] paired with the reduced partition expression it produces.
+ *
+ * When a key-grouped partitioning is reduced onto another partitioning's key space, the original
+ * partition expressions no longer describe the reduced keys (their data type and their value are
+ * those of the target key space). This pair carries both the reducer and the expression the
+ * reduced keys correspond to, so the output partitioning can report a type-correct expression.
+ *
+ * @param reducer reducer that maps this side's partition key values onto the other side's key space
+ * @param reducedExpression expression whose data type matches the reduced keys. Stored as-is from
+ * the expression pair that produced the reducer; the only structural
+ * consumer, `GroupPartitionsExec`, re-targets it at each reported
+ * `KeyedPartitioning`'s own key attribute in `outputPartitioning` and
+ * normalizes it positionally in `doCanonicalize`, so the attribute it
+ * carries here is not load-bearing.
+ */
+case class KeyReducer(reducer: Reducer[_, _], reducedExpression: TransformExpression)
+
+/**
+ * A [[Reducer]] that reduces an identity transform onto a partition transform: it applies the
+ * given partition transform to the raw value of the identity transform's key. For instance, an
+ * `identity(id)` side joined to a `bucket(4, id)` side reduces with
+ * `IdentityReducer(bucket(4, id))`, which maps each raw id value to its `bucket(4, id)` value.
+ *
+ * It is a case class so that structurally identical reducers compare by value, and plans holding
+ * them stay canonicalization-equal.
+ *
+ * @param transform the partition transform, re-targeted at the identity side's key attribute
+ */
+case class IdentityReducer(transform: TransformExpression) extends Reducer[Any, Any] {
+ // `transform` has a single leaf attribute (`KeyedPartitioning.supportsExpressions`), which is
+ // bound to ordinal 0 of the single-value row `reduce` evaluates it against.
+ @transient private lazy val bound: Expression =
+ BindReferences.bindReference(transform, AttributeSeq(transform.references.toSeq))
+
+ override def reduce(v: Any): Any = bound.eval(new GenericInternalRow(Array[Any](v)))
+
+ override def resultType(): DataType = transform.dataType
+
+ override def displayName(): String = transform.toString
+}
+
/**
* [[ShuffleSpec]] created by [[KeyedPartitioning]].
*
@@ -1066,44 +1108,82 @@ case class KeyedShuffleSpec(
}
/**
- * Return a set of [[Reducer]] for the partition expressions of this shuffle spec,
- * on the partition expressions of another shuffle spec.
- * <p>
- * A [[Reducer]] exists for a partition expression function of this shuffle spec if it is
- * 'reducible' on the corresponding partition expression function of the other shuffle spec.
+ * Compute the reducers for both sides of a join between this shuffle spec and `other`, in a
+ * single pass over the two sides' partition expressions. A pair's reducer lookups
+ * (`TransformExpression.reducers`) are shared between the two directions: the reverse lookup a
+ * direction needs to detect a single-side reduce is exactly the other direction's reducer, so
+ * this materializes each catalog `Reducer` once per direction.
* <p>
- * If a value is returned, there must be one [[Reducer]] per partition expression.
- * A None value in the set indicates that the particular partition expression is not reducible
- * on the corresponding expression on the other shuffle spec.
+ * A [[Reducer]] exists for a partition expression of one side if it is 'reducible' on the
+ * corresponding partition expression of the other side. If a side's value is returned, there
+ * must be one entry per partition expression of that side. A None entry indicates that the
+ * particular partition expression is not reducible on the corresponding expression.
* <p>
- * Returning none also indicates that none of the partition expressions can be reduced on the
- * corresponding expression on the other shuffle spec.
+ * Returning none for a side indicates that none of its partition expressions can be reduced on
+ * the corresponding expression of the other side.
*
* @param other other key-grouped shuffle spec
+ * @return the reducers for this side and for `other` w.r.t. each other
*/
- def reducers(other: KeyedShuffleSpec): Option[Seq[Option[Reducer[_, _]]]] = {
- val results = partitioning.expressions.zip(other.partitioning.expressions).map {
- case (e1: TransformExpression, e2: TransformExpression) => e1.reducers(e2)
+ def reducersBothWays(other: KeyedShuffleSpec)
+ : (Option[Seq[Option[KeyReducer]]], Option[Seq[Option[KeyReducer]]]) = {
+ val results: Seq[(Option[KeyReducer], Option[KeyReducer])] =
+ partitioning.expressions.zip(other.partitioning.expressions).map {
+ case (e1: TransformExpression, e2: TransformExpression) =>
+ val thisReducer = e1.reducers(e2)
+ val otherReducer = e2.reducers(e1)
+
+ val thisResult = thisReducer.map { reducer =>
+ if (otherReducer.isEmpty) {
+ // Only this side reduces. The reducer contract is r(f1(x)) = f2(x) where "=" matches
+ // both value and data type, so the reduced keys equal the target transform applied to
+ // this side's child. Report the target transform (re-targeted at the reporting
+ // partitioning's key attribute by `GroupPartitionsExec`) instead of the un-reduced
+ // `e1`, whose type can differ from the reduced keys. A connector that violates the
+ // contract with a reducer of a different result type fails the reduced-types check in
+ // `EnsureRequirements`, since the other side's keys are typed by the target transform.
+ KeyReducer(reducer, e2)
+ } else {
+ // Both sides reduce: the reduced keys are r1(f1(x)) = r2(f2(x)), which no single
+ // transform describes. Keep reporting the original expression; its type can differ
+ // from the reduced keys, which a downstream shuffle or grouping can then fail on with
+ // a ClassCastException. Known gap, tracked in SPARK-59121.
+ KeyReducer(reducer, e1)
+ }
+ }
+
+ val otherResult = otherReducer.map { reducer =>
+ if (thisReducer.isEmpty) {
+ // Only the other side reduces; symmetric to the single-side case above.
+ KeyReducer(reducer, e1)
+ } else {
+ // Both sides reduce; symmetric to the both-sides case above.
+ KeyReducer(reducer, e2)
+ }
+ }
+
+ (thisResult, otherResult)
// Identity transform on this side, arbitrary transform on the other side: create a reducer
- // that applies the other's transform to the raw identity values. The symmetric case
- // (TransformExpression, AttributeReference) is handled when the other side calls reducers.
- // Each partition expression is guaranteed to have exactly one leaf child (asserted in
- // keyPositions), so `a` lives at position 0 in the row we construct.
+ // that applies the other's transform to the raw identity values. Each partition expression
+ // is guaranteed to have exactly one leaf child (asserted in keyPositions), which
+ // `IdentityReducer` binds to ordinal 0.
case (a: AttributeReference, t: TransformExpression) =>
- val reducerExpr = t.transform { case _: AttributeReference => a }
- val boundExpr = BindReferences.bindReference(reducerExpr, AttributeSeq(Seq(a)))
- Some(new Reducer[Any, Any] {
- override def reduce(v: Any): Any = boundExpr.eval(new GenericInternalRow(Array[Any](v)))
- override def resultType(): DataType = reducerExpr.dataType
- override def displayName(): String = reducerExpr.toString
- })
-
- case (_, _) => None
+ (Some(KeyReducer(IdentityReducer(t.withReference(a)), t)), None)
+
+ // Symmetric: identity transform on the other side.
+ case (t: TransformExpression, a: AttributeReference) =>
+ (None, Some(KeyReducer(IdentityReducer(t.withReference(a)), t)))
+
+ case (_, _) => (None, None)
}
// optimize to not return a value, if none of the partition expressions are reducible
- if (results.forall(p => p.isEmpty)) None else Some(results)
+ val thisReducers = results.map(_._1)
+ val otherReducers = results.map(_._2)
+ val thisResult = if (thisReducers.forall(_.isEmpty)) None else Some(thisReducers)
+ val otherResult = if (otherReducers.forall(_.isEmpty)) None else Some(otherReducers)
+ (thisResult, otherResult)
}
override def canCreatePartitioning: Boolean =
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala
index dc07c2a59561..6e11c0fe5e4b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala
@@ -43,11 +43,11 @@ import org.apache.spark.sql.catalyst.parser.ParseException
import org.apache.spark.sql.catalyst.plans.JoinType
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.ValueInterval
+import org.apache.spark.sql.catalyst.plans.physical.KeyReducer
import org.apache.spark.sql.catalyst.trees.{Origin, TreeNode}
import org.apache.spark.sql.catalyst.util.{sideBySide, CharsetProvider, DateTimeUtils, FailFastMode, IntervalUtils, MapData}
import org.apache.spark.sql.connector.catalog.{CatalogNotFoundException, Table, TableProvider}
import org.apache.spark.sql.connector.catalog.CatalogV2Implicits._
-import org.apache.spark.sql.connector.catalog.functions.Reducer
import org.apache.spark.sql.connector.expressions.Transform
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.internal.StaticSQLConf.GLOBAL_TEMP_DATABASE
@@ -3174,12 +3174,12 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE
}
def storagePartitionJoinIncompatibleReducedTypesError(
- leftReducers: Option[Seq[Option[Reducer[_, _]]]],
+ leftReducers: Option[Seq[Option[KeyReducer]]],
leftReducedDataTypes: Seq[DataType],
- rightReducers: Option[Seq[Option[Reducer[_, _]]]],
+ rightReducers: Option[Seq[Option[KeyReducer]]],
rightReducedDataTypes: Seq[DataType]): Throwable = {
- def reducersNames(reducers: Option[Seq[Option[Reducer[_, _]]]]) = {
- reducers.toSeq.flatMap(_.map(_.map(_.displayName()).getOrElse("identity")))
+ def reducersNames(reducers: Option[Seq[Option[KeyReducer]]]) = {
+ reducers.toSeq.flatMap(_.map(_.map(_.reducer.displayName()).getOrElse("identity")))
.mkString("[", ", ", "]")
}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala
index 494fb38a305e..155e7d2cc305 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala
@@ -25,9 +25,9 @@ import org.apache.spark.rdd.{CoalescedRDD, PartitionCoalescer, PartitionGroup, R
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.expressions.codegen.GenerateOrdering
-import org.apache.spark.sql.catalyst.plans.physical.{KeyedPartitioning, Partitioning}
+import org.apache.spark.sql.catalyst.plans.QueryPlan
+import org.apache.spark.sql.catalyst.plans.physical.{IdentityReducer, KeyedPartitioning, KeyReducer, Partitioning}
import org.apache.spark.sql.catalyst.util.{truncatedString, InternalRowComparableWrapper}
-import org.apache.spark.sql.connector.catalog.functions.Reducer
import org.apache.spark.sql.execution.{SafeForKWayMerge, SparkPlan, UnaryExecNode}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.DataType
@@ -59,7 +59,7 @@ case class GroupPartitionsExec(
child: SparkPlan,
@transient joinKeyPositions: Option[Seq[Int]] = None,
@transient expectedPartitionKeys: Option[Seq[(InternalRowComparableWrapper, Int)]] = None,
- @transient reducers: Option[Seq[Option[Reducer[_, _]]]] = None,
+ @transient reducers: Option[Seq[Option[KeyReducer]]] = None,
@transient distributePartitions: Boolean = false,
@transient enableSortedMerge: Boolean = false
) extends UnaryExecNode {
@@ -67,10 +67,13 @@ case class GroupPartitionsExec(
override def outputPartitioning: Partitioning = {
child.outputPartitioning match {
case p: Partitioning with Expression =>
- // There can be multiple `KeyedPartitioning` in an output partitioning of a join, but they
- // can only differ in `expressions`. `partitionKeys` must match so we can calculate it only
- // once via `groupedPartitions`.
-
+ // There can be multiple `KeyedPartitioning`s in an output partitioning of a join, but they
+ // can only differ in `expressions`. `partitionKeys` must match, so they are calculated
+ // only once via `groupedPartitions`. When reducers are applied, the stored reduced
+ // expressions are re-targeted at each `KeyedPartitioning`'s own key attribute and reported
+ // instead of the original ones. Their data types match the reduced partition keys for the
+ // identity-vs-transform and single-side-transform reducers; for the both-sides-reduce
+ // shape no single transform describes the keys (see `KeyedShuffleSpec.reducersBothWays`).
val keyedPartitionings = p.collect { case k: KeyedPartitioning => k }
if (keyedPartitionings.size > 1) {
val first = keyedPartitionings.head
@@ -83,13 +86,51 @@ case class GroupPartitionsExec(
p.transform {
case k: KeyedPartitioning =>
val projectedExpressions = joinKeyPositions.fold(k.expressions)(_.map(k.expressions))
- KeyedPartitioning(projectedExpressions, groupedPartitions.map(_._1),
+ val effectiveExpressions = reducers match {
+ case Some(exprs) =>
+ assert(projectedExpressions.length == exprs.length)
+ projectedExpressions.zip(exprs).map {
+ case (expr, Some(KeyReducer(_, reduced))) =>
+ // `reduced` was stored from the single spec that `createKeyedShuffleSpec`
+ // picked (`collectFirst`); re-target it at this `KeyedPartitioning`'s own key
+ // attribute so that every `KeyedPartitioning` in a collection keeps its own.
+ reduced.withReference(expr.references.head)
+ case (expr, None) => expr
+ }
+ case None => projectedExpressions
+ }
+ KeyedPartitioning(effectiveExpressions, groupedPartitions.map(_._1),
isGrouped = isGrouped)
}.asInstanceOf[Partitioning]
case o => o
}
}
+ override def doCanonicalize(): SparkPlan = {
+ // `KeyReducer` is a plain case class, not an `Expression`, so plan canonicalization does not
+ // normalize the exprIds inside it. Normalize them positionally here: `reducedExpression` may
+ // reference the other join side's key, which this node's child does not output, so both it
+ // and the identity reducer's transform are normalized against their own references. Without
+ // this, structurally identical SPJ subtrees with value-equal reducers stop comparing equal
+ // once reducers are applied, and exchange/subquery reuse silently stops deduplicating them.
+ val canonicalized = super.doCanonicalize().asInstanceOf[GroupPartitionsExec]
+ canonicalized.copy(reducers = canonicalized.reducers.map { reducerSeqs =>
+ reducerSeqs.map { keyReducerOpt =>
+ keyReducerOpt.map { keyReducer =>
+ val reducer = keyReducer.reducer match {
+ case r: IdentityReducer => r.copy(transform = QueryPlan.normalizeExpressions(
+ r.transform, r.transform.references.toSeq).asInstanceOf[TransformExpression])
+ case other => other
+ }
+ keyReducer.copy(
+ reducer = reducer,
+ reducedExpression = QueryPlan.normalizeExpressions(keyReducer.reducedExpression,
+ keyReducer.reducedExpression.references.toSeq).asInstanceOf[TransformExpression])
+ }
+ }
+ })
+ }
+
/**
* Aligns partitions based on `expectedPartitionKeys` and clustering mode.
*/
@@ -334,7 +375,7 @@ case class GroupPartitionsExec(
}.iterator
val expectedStr = expectedPartitionKeys.map(ks => s"ExpectedPartitionKeys: ${ks.size}")
val reducersStr = reducers.map { seq =>
- val names = seq.map(_.map(_.displayName()).getOrElse("identity"))
+ val names = seq.map(_.map(_.reducer.displayName()).getOrElse("identity"))
s"Reducers: ${truncatedString(names, "[", ", ", "]", joinKeyMaxFields)}"
}
val distributeStr = Iterator(s"DistributePartitions: $distributePartitions")
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
index b5114aa0be90..1b27239e4c04 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
@@ -27,7 +27,6 @@ import org.apache.spark.sql.catalyst.plans._
import org.apache.spark.sql.catalyst.plans.physical._
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper
-import org.apache.spark.sql.connector.catalog.functions.Reducer
import org.apache.spark.sql.errors.QueryExecutionErrors
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.datasources.v2.GroupPartitionsExec
@@ -549,8 +548,7 @@ case class EnsureRequirements(
// in case of compatible but not identical partition expressions, we apply 'reduce'
// transforms to group one side's partitions as well as the common partition values
- val leftReducers = leftSpec.reducers(rightSpec)
- val rightReducers = rightSpec.reducers(leftSpec)
+ val (leftReducers, rightReducers) = leftSpec.reducersBothWays(rightSpec)
val (leftReducedDataTypes, leftReducedKeys) = leftReducers.fold(
(leftPartitioning.keyDataTypes, leftPartitioning.partitionKeys)
)(leftPartitioning.reduceKeys)
@@ -742,7 +740,7 @@ case class EnsureRequirements(
plan: SparkPlan,
joinKeyPositions: Option[Seq[Int]],
mergedPartitionKeys: Seq[(InternalRowComparableWrapper, Int)],
- reducers: Option[Seq[Option[Reducer[_, _]]]],
+ reducers: Option[Seq[Option[KeyReducer]]],
distributePartitions: Boolean): SparkPlan = {
plan match {
case g: GroupPartitionsExec =>
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
index e9dc7ee56e06..9d3d5d64a1bc 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala
@@ -23,14 +23,21 @@ import org.apache.spark.{SparkConf, SparkException}
import org.apache.spark.rdd.SortedMergeCoalescedRDD
import org.apache.spark.sql.{DataFrame, ExplainSuiteHelper, Row}
import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, Literal, TransformExpression}
+import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, ExprId, Literal, TransformExpression}
import org.apache.spark.sql.catalyst.plans.physical
import org.apache.spark.sql.connector.catalog.{Column, Identifier, InMemoryTableCatalog}
import org.apache.spark.sql.connector.catalog.functions._
import org.apache.spark.sql.connector.distributions.Distributions
import org.apache.spark.sql.connector.expressions._
import org.apache.spark.sql.connector.expressions.Expressions._
-import org.apache.spark.sql.execution.{ExtendedMode, FormattedMode, RDDScanExec, SimpleMode, SortExec, SparkPlan}
+import org.apache.spark.sql.execution.{
+ ExtendedMode,
+ FormattedMode,
+ LocalTableScanExec,
+ RDDScanExec,
+ SimpleMode,
+ SortExec,
+ SparkPlan}
import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation, GroupPartitionsExec}
import org.apache.spark.sql.execution.exchange.{ShuffleExchangeExec, ShuffleExchangeLike}
import org.apache.spark.sql.execution.joins.{ShuffledHashJoinExec, SortMergeJoinExec}
@@ -533,6 +540,168 @@ 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,
+ // so downstream consumers see the keys' real type: another join can reduce onto the reported
+ // transform, and a reduced layout can serve as another child's shuffle target only when the
+ // expressions describe the keys.
+ 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_JOIN_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.
+ // This test deliberately leaves `V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS` off:
+ // both joins use the whole partition key, and without the config the failure on base is
+ // exercised through the second, independent trigger (`reduceKeys` at the second join) instead
+ // of `createShuffleSpec` -> `toGrouped`.
+ 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") {
+ checkAnswer(df, Seq(
+ Row(1, "a", "x", "p"), Row(2, "b", "y", "q"), Row(3, "c", "z", "r")))
+ assert(collectShuffles(df.queryExecution.executedPlan).isEmpty,
+ "storage-partitioned join should not shuffle")
+ }
+ }
+
+ test("SPARK-59045: reduced expression is retargeted per KeyedPartitioning") {
+ // A chained SPJ's output partitioning reports one `KeyedPartitioning` per join side, but the
+ // reduced expression is derived from the single spec that `createKeyedShuffleSpec` picks
+ // (`collectFirst`). Re-targeting it at each `KeyedPartitioning`'s own key attribute keeps the
+ // other sides' partitionings intact - otherwise a GROUP BY on the other side's key no longer
+ // sees a partitioning on it and the query shuffles (0 shuffles on base and here, 1 if the
+ // use-site re-targeting is dropped).
+ val cols = Array(Column.create("id", LongType), Column.create("data", StringType))
+ createTable("b16", cols, Array(bucket(16, "id")))
+ createTable("b8", cols, Array(bucket(8, "id")))
+ createTable("b4", cols, Array(bucket(4, "id")))
+ val values = (0 until 16).map(i => s"($i, 'v$i')").mkString(", ")
+ Seq("b16", "b8", "b4").foreach(t => sql(s"INSERT INTO testcat.ns.$t VALUES $values"))
+
+ val df = sql(
+ "SELECT b8.id, count(*) FROM testcat.ns.b16 " +
+ "JOIN testcat.ns.b8 ON b16.id = b8.id JOIN testcat.ns.b4 ON b16.id = b4.id " +
+ "GROUP BY b8.id")
+
+ withSQLConf(SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ checkAnswer(df, (0 until 16).map(i => Row(i.toLong, 1L)))
+ assert(collectShuffles(df.queryExecution.executedPlan).isEmpty,
+ "storage-partitioned join should not shuffle")
+ }
+ }
+
+ test("SPARK-59045: compatible transforms reduce data type with subset join keys") {
+ // The join is on `id`, a subset of the partition keys `[identity(dt), identity(id)]` and
+ // `[identity(dt), bucket(2, id)]`. The identity(id) side is reduced to bucket(2, id), whose
+ // data type differs, while the dt partition key is projected away.
+ val cols = Array(
+ Column.create("id", LongType),
+ Column.create("dt", StringType),
+ Column.create("data", StringType))
+ createTable("t1", cols, Array(identity("dt"), identity("id")))
+ createTable("t2", cols, Array(identity("dt"), bucket(2, "id")))
+ sql("INSERT INTO testcat.ns.t1 VALUES (1, '2020', 'a'), (2, '2020', 'b'), (3, '2021', 'c')")
+ sql("INSERT INTO testcat.ns.t2 VALUES (1, '2020', 'x'), (2, '2020', 'y'), (3, '2021', 'z')")
+
+ val df = sql(
+ "SELECT t1.id, t1.dt, t1.data, t2.dt, t2.data FROM testcat.ns.t1 " +
+ "JOIN testcat.ns.t2 ON t1.id = t2.id")
+
+ withSQLConf(
+ SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION.key -> "false",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_JOIN_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") {
+ checkAnswer(df, Seq(
+ Row(1, "2020", "a", "2020", "x"), Row(2, "2020", "b", "2020", "y"),
+ Row(3, "2021", "c", "2021", "z")))
+ assert(collectShuffles(df.queryExecution.executedPlan).isEmpty,
+ "storage-partitioned join should not shuffle")
+ }
+ }
+
+ test("SPARK-59045: canonicalization normalizes the reduced expressions") {
+ // `KeyReducer` is a plain case class, not an `Expression`, so plan canonicalization does not
+ // normalize the exprIds inside it. Two structurally identical `GroupPartitionsExec`s with
+ // value-equal reducers must still compare equal after canonicalization, or exchange/subquery
+ // reuse silently stops deduplicating their subtrees. The reduced expression is the other join
+ // side's transform, so it references an attribute this node's child does not output; the
+ // identity reducer's transform references this side's own key.
+ val a1 = AttributeReference("id", LongType, nullable = true)().withExprId(ExprId(1))
+ val a2 = AttributeReference("id", LongType, nullable = true)().withExprId(ExprId(2))
+ val o1 = AttributeReference("oid", LongType, nullable = true)().withExprId(ExprId(11))
+ val o2 = AttributeReference("oid", LongType, nullable = true)().withExprId(ExprId(12))
+ def groupPartitions(
+ attr: AttributeReference,
+ otherAttr: AttributeReference,
+ reducer: Reducer[_, _]): GroupPartitionsExec = {
+ val child = new LocalTableScanExec(Seq(attr), Nil, None)
+ val reduced = TransformExpression(BucketFunction, Seq(otherAttr), Some(2))
+ GroupPartitionsExec(child,
+ reducers = Some(Seq(Some(physical.KeyReducer(reducer, reduced)))))
+ }
+ // A value-equal reducer, with the reduced expression over the other side's attribute.
+ assert(groupPartitions(a1, o1, BucketReducer(2)).canonicalized ==
+ groupPartitions(a2, o2, BucketReducer(2)).canonicalized)
+ // The identity-derived reducer, whose transform is over this side's own attribute.
+ assert(groupPartitions(a1, o1, physical.IdentityReducer(
+ TransformExpression(BucketFunction, Seq(a1), Some(2)))).canonicalized ==
+ groupPartitions(a2, o2, physical.IdentityReducer(
+ TransformExpression(BucketFunction, Seq(a2), Some(2)))).canonicalized)
+ // Structurally different reducers stay unequal after canonicalization.
+ assert(groupPartitions(a1, o1, BucketReducer(2)).canonicalized !=
+ groupPartitions(a2, o2, BucketReducer(3)).canonicalized)
+
+ // A multi-key partitioning with a mixed reducer sequence (identity keys are not reducible):
+ // the normalization is per position, and the None entries pass through untouched.
+ val dt1 = AttributeReference("dt", StringType, nullable = true)().withExprId(ExprId(21))
+ val dt2 = AttributeReference("dt", StringType, nullable = true)().withExprId(ExprId(22))
+ def mixedKeyGroupPartitions(
+ attr: AttributeReference,
+ dt: AttributeReference,
+ otherAttr: AttributeReference): GroupPartitionsExec = {
+ val child = new LocalTableScanExec(Seq(attr, dt), Nil, None)
+ val reduced = TransformExpression(BucketFunction, Seq(otherAttr), Some(2))
+ GroupPartitionsExec(child,
+ reducers = Some(Seq(None, Some(physical.KeyReducer(BucketReducer(2), reduced)))))
+ }
+ assert(mixedKeyGroupPartitions(a1, dt1, o1).canonicalized ==
+ mixedKeyGroupPartitions(a2, dt2, o2).canonicalized)
+ }
+
test("partitioned join: join with two partition keys and matching & sorted partitions") {
val items_partitions = Array(bucket(8, "id"), days("arrive_time"))
createTable(items, itemsColumns, items_partitions)
@@ -3694,18 +3863,29 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
s"(5, 44.0, cast('2020-01-15' as timestamp)), " +
s"(7, 46.5, cast('2021-02-08' as timestamp))")
+ // A third table partitioned by `identity(time)` joins on the same timestamps: its side
+ // reduces onto the first join's reported `years(arrive_time)`, so the chain plans without a
+ // shuffle only while the reduced keys are reported under the type-correct target transform.
+ val shipments = "shipments"
+ createTable(shipments, purchasesColumns, Array(identity("time")))
+ sql(s"INSERT INTO testcat.ns.$shipments VALUES " +
+ s"(1, 42.0, cast('2020-01-01' as timestamp)), " +
+ s"(9, 46.5, cast('2021-02-08' as timestamp))")
+
withSQLConf(
SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
Seq(
- s"testcat.ns.$items i JOIN testcat.ns.$purchases p ON p.time = i.arrive_time",
- s"testcat.ns.$purchases p JOIN testcat.ns.$items i ON i.arrive_time = p.time"
+ s"testcat.ns.$items i JOIN testcat.ns.$purchases p ON p.time = i.arrive_time " +
+ s"JOIN testcat.ns.$shipments s ON i.arrive_time = s.time",
+ s"testcat.ns.$purchases p JOIN testcat.ns.$items i ON i.arrive_time = p.time " +
+ s"JOIN testcat.ns.$shipments s ON i.arrive_time = s.time"
).foreach { joinString =>
val df = sql(
s"""
- |${selectWithMergeJoinHint("i", "p")} id, item_id
+ |${selectWithMergeJoinHint("i", "p")} i.id, p.item_id
|FROM $joinString
- |ORDER BY id, item_id
+ |ORDER BY i.id, p.item_id
|""".stripMargin)
val shuffles = collectShuffles(df.queryExecution.executedPlan)
@@ -4203,11 +4383,10 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
s"(2, 'bb', cast('2021-06-01' as timestamp))")
}
- test("SPARK-59120: incompatible reduced key types are reported instead of cast") {
- // The first join reduces the identity side onto the year key space. The third table does not
- // reduce at all, so the two key spaces genuinely differ, and reading each side's keys at their
- // own types is what lets `EnsureRequirements` say so. On `master` the merge cast one to the
- // other and threw `ClassCastException`.
+ test("SPARK-59120: a second join with a non-reducing side plans on the reduced keys") {
+ // The first join reduces the identity side onto the year key space and reports the reduced
+ // expression `years(a.ts)`, so the second join's identity side reduces onto it as well and
+ // the whole query plans without a shuffle.
withTable("t_identity", "t_years", "t_identity2") {
createTsTable("t_identity", Array(identity("ts")))
createTsTable("t_years", Array(years("ts")))
@@ -4223,22 +4402,18 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
- checkError(
- exception = intercept[SparkException](sql(query).collect()),
- condition = "STORAGE_PARTITION_JOIN_INCOMPATIBLE_REDUCED_TYPES",
- parameters = Map(
- "leftReducers" -> "[]",
- "leftReducedDataTypes" -> "[\"INT\"]",
- "rightReducers" -> "[]",
- "rightReducedDataTypes" -> "[\"TIMESTAMP\"]"))
+ val df = sql(query)
+ checkAnswer(df, Seq(Row(1, "aa"), Row(2, "bb")))
+ assert(collectAllShuffles(df.queryExecution.executedPlan).isEmpty,
+ "the second join reduces onto the type-correct reported expression")
}
}
}
- test("SPARK-59120: another child is not shuffled onto reducer-rewritten keys") {
- // `canCreatePartitioning` refuses the reduced side as a layout for the second join, so both of
- // that join's sides are shuffled. Without the gate the driver dies while assembling that
- // shuffle's key map, where the stored keys are re-wrapped at the expressions' types.
+ test("SPARK-59120: another child is shuffled onto the type-correct reduced keys") {
+ // The reduced side's expression is reported as `years(a.ts)`, which describes the reduced
+ // keys, so `canCreatePartitioning`'s shape gate accepts the reduced layout and only the
+ // unpartitioned side is shuffled onto it.
withTable("t_identity", "t_years", "t_plain") {
createTsTable("t_identity", Array(identity("ts")))
createTsTable("t_years", Array(years("ts")))
@@ -4257,8 +4432,8 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with
SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
val df = sql(query)
checkAnswer(df, Seq(Row(1, "aa"), Row(2, "bb")))
- assert(collectAllShuffles(df.queryExecution.executedPlan).size == 2,
- "the second join shuffles both of its sides")
+ assert(collectAllShuffles(df.queryExecution.executedPlan).size == 1,
+ "the second join shuffles only the unpartitioned side, onto the years(ts) layout")
}
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org