dongjoon-hyun commented on code in PR #58345:
URL: https://github.com/apache/spark/pull/58345#discussion_r3874270256
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala:
##########
@@ -463,8 +465,22 @@ object ShuffleExchangeExec {
val projection =
UnsafeProjection.create(sortingExpressions.map(_.child), outputAttributes)
row => projection(row)
case SinglePartition => identity
- case KeyedPartitioning(expressions, _, _, _) =>
- row => bindReferences(expressions, outputAttributes).map(_.eval(row))
+ case k: KeyedPartitioning =>
+ val expressions = bindReferences(k.expressions,
outputAttributes).toArray
+ // Wrap the evaluated partition key so it compares equal to the
KeyGroupedPartitioner's
+ // map keys (the partitioning's own wrappers) under `RowOrdering`
semantics. The wrapped
+ // row is reused across records; KeyGroupedPartitioner does not retain
lookup keys.
+ val wrapperFactory = InternalRowComparableWrapper
+ .getInternalRowComparableWrapperFactory(k.expressionDataTypes)
Review Comment:
Good catch, @sunchao — confirmed. I reproduced it with a regression test
joining an identity-partitioned `STRUCT<a:INT>` column with an unpartitioned
`STRUCT<b:INT>` column: no cast is inserted, the one-sided shuffle stays
eligible, and matches were silently lost because
`InternalRowComparableWrapper.equals` rejects the differing `dataTypes` before
comparing values.
Fixed in a40ee40: the driver side now re-wraps `partitionKeys` rows through
the same `getInternalRowComparableWrapperFactory(k.expressionDataTypes)` the
executor-side lookup uses, so the stored keys and the lookup keys always share
one comparison schema regardless of which side's wrappers
`KeyedShuffleSpec.createPartitioning` retained.
Added the regression test you suggested (`SPARK-59054: shuffle one side:
struct partition keys with different field names`, multiple keyed partitions,
both conf arms compared); verified it fails before this change and passes
after. The suite is 107/107.
--
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]
