dongjoon-hyun commented on code in PR #58335:
URL: https://github.com/apache/spark/pull/58335#discussion_r3876091918
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -713,17 +714,18 @@ 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()
Review Comment:
**[design]** The reduced key type now has two unvalidated sources of truth:
`reduceKeys` uses `reducer.resultType()` here, while
`GroupPartitionsExec.outputPartitioning` reports `reducedExpression.dataType`
(the target function's `resultType()`). `Reducer`'s contract implies they
agree, but nothing validates it, so a connector whose `Reducer.resultType()`
disagrees with the target transform's type recreates the exact
expressions-vs-keys mismatch this PR fixes, hidden inside the new mechanism.
Deriving the reduced types from `reducedExpression.dataType`, or asserting
agreement when constructing `KeyReducer`, would close the drift by construction.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -59,7 +58,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,
Review Comment:
**[correctness]** `KeyReducer` embeds an exprId-bearing
`TransformExpression` in a non-`Expression` constructor arg.
`QueryPlan.doCanonicalize` normalizes exprIds only via `mapExpressions`, which
does not recurse into the plain `KeyReducer` case class (`@transient` affects
serialization only, not `equals`). Pre-PR, a connector returning value-equal
`Reducer` instances on the `ReducibleFunction` path (e.g. the test
`BucketReducer(divisor)` case class) allowed two identical SPJ subtrees'
`GroupPartitionsExec` nodes to compare equal after canonicalization; post-PR
the side-specific `AttributeReference` exprIds inside `reducedExpression` break
that equality, so exchange/subquery/stage reuse can silently stop deduplicating
those subtrees.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1374,9 +1391,25 @@ 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[KeyReducer]]] = {
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 =>
+ if (e2.reducers(e1).isEmpty) {
Review Comment:
**[efficiency]** `e2.reducers(e1)` is invoked per expression pair solely to
test `.isEmpty`, materializing a reverse `Reducer` via the catalog function's
`reducer()` call. `TransformExpression.reducers` is an uncached `def`, and
`EnsureRequirements` already calls `spec.reducers` in both directions, so each
direction's catalog lookup now runs twice per join-key pair (4 instead of 2) on
every SPJ planning pass, including AQE re-planning. A slow or allocating
third-party `ReducibleFunction.reducer()` pays this on a discarded probe;
computing both directions once per pair would avoid it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -70,11 +69,29 @@ 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 are reported
instead of the original
+ // ones. For the identity-vs-transform and single-side-transform
reducers their data type
+ // matches the reduced partition keys by construction; for the
both-sides-reduce shape no
+ // single transform describes the keys (see
`KeyedShuffleSpec.reducers`).
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(KeyReducer(_, reduced))) =>
+ // `reduced` was derived 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.
+ val attr = expr.references.head
+ reduced.transform { case _: AttributeReference => attr }
Review Comment:
**[reuse]** This is the third copy of the `X.transform { case _:
AttributeReference => attr }` retargeting idiom (also at
`KeyedShuffleSpec.reducers` in both branches), spread across two modules with
no shared helper. This copy must stay exactly consistent with how
`reducedExpression` was built in `KeyedShuffleSpec.reducers` - a future change
to the rebind rule (e.g. per-position rebinding for multi-reference or
nested-field transforms) can silently miss one copy, making the reported
`outputPartitioning` stop matching the actual reduced keys. A small helper on
`TransformExpression` (returning `TransformExpression`, which also removes the
`asInstanceOf` casts) would fit this PR.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -567,7 +567,8 @@ 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]]):
Review Comment:
**[style]** Nit: this signature (and `object KeyedPartitioning.reduceKeys`
below) was re-wrapped with the bare return type on its own continuation line,
but the one-line forms fit within the 100-char limit after the type shortened
(`Seq[Option[KeyReducer]]` is shorter than `Seq[Option[Reducer[_, _]]]`), and
this wrapping style is not used elsewhere in the file. Keeping the original
one-line form avoids the formatting churn.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1374,9 +1391,25 @@ 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[KeyReducer]]] = {
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 =>
+ if (e2.reducers(e1).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 that expression (re-targeted at this
side's attribute)
+ // instead of the un-reduced `e1`, whose type can differ from the
reduced keys.
+ val thisSideChild = e1.references.head
+ val reducedExpr = e2.transform { case _: AttributeReference =>
thisSideChild }
+ KeyReducer(reducer, reducedExpr.asInstanceOf[TransformExpression])
+ } else {
+ // Both sides reduce: the reduced keys are r1(f1(x)) = r2(f2(x)),
which no single
+ // transform describes. Keep reporting the original expression.
Known gap, tracked in
+ // the follow-up for SPARK-59045.
+ KeyReducer(reducer, e1)
Review Comment:
**[correctness]** The both-sides-reduce branch keeps the un-reduced `e1`, so
the `ClassCastException` this PR fixes remains reachable in that shape with no
runtime fail-safe. The shape is reachable even with in-tree test functions
(`BucketFunction`'s gcd reducer reduces both ways;
`DaysFunctionWithToYearsReducerWithLongResult` /
`YearsFunctionWithToYearsReducerWithLongResult` both reduce to `LongType` while
their transforms report `DateType`/`IntegerType`). `EnsureRequirements`'
`storagePartitionJoinIncompatibleReducedTypesError` check passes here because
the LEFT and RIGHT *reduced* types equal each other, so `GroupPartitionsExec`
reports original-typed expressions over reduced-typed partition keys, and a
downstream GROUP BY / second join / shuffle deriving an ordering from
`expressionDataTypes` throws the same CCE.
Since the follow-up is tracked separately, could this branch fail safe until
then, e.g. return `None` (fall back to shuffle) or raise the dedicated error
when `reducer.resultType() != e1.dataType`?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1374,9 +1391,25 @@ 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[KeyReducer]]] = {
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 =>
+ if (e2.reducers(e1).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 that expression (re-targeted at this
side's attribute)
+ // instead of the un-reduced `e1`, whose type can differ from the
reduced keys.
+ val thisSideChild = e1.references.head
+ val reducedExpr = e2.transform { case _: AttributeReference =>
thisSideChild }
Review Comment:
**[simplification]** This creation-time retargeting is dead work: the only
structural consumer of `reducedExpression`
(`GroupPartitionsExec.outputPartitioning`) unconditionally retargets it again
at each `KeyedPartitioning`'s own attribute, and all other reads use only
`.reducer` (`reduceKeys`, `displayName` in `QueryExecutionErrors`). Storing raw
`e2` and doing the single retarget at the use site is equivalent, removes one
tree transform + allocation per expression pair, and stops implying that the
attribute chosen here (from whichever spec `collectFirst` picked) is
load-bearing.
--
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]
