dongjoon-hyun opened a new pull request, #58345: URL: https://github.com/apache/spark/pull/58345
### What changes were proposed in this pull request? This PR fixes a correctness bug in the `KeyedPartitioning` shuffle path (storage-partitioned join with `spark.sql.sources.v2.bucketing.shuffle.enabled=true`) by making the partition-key lookup in `KeyGroupedPartitioner` compare keys by value. - `ShuffleExchangeExec.getPartitioner`: the driver-side `valueMap` keys are now `UnsafeRow`s produced by an `UnsafeProjection` over the partition key data types, instead of `Seq[Any]` from `InternalRow.toSeq`. - `ShuffleExchangeExec.getPartitionKeyExtractor`: the executor-side lookup key is built with the same `UnsafeProjection` (the bound partition expressions are evaluated into a reused `GenericInternalRow` and then projected), instead of `Seq[Any]` from per-row `eval`. Since `UnsafeProjection` cannot generate code for `TransformExpression`, the expressions are evaluated first and the projection only reads the resulting row. - `KeyGroupedPartitioner` (core): takes `Map[Any, Int]` and looks keys up with `getOrElse`, removing the `ArraySeq` normalization and the `getOrElseUpdate` map mutation. Because lookup keys are no longer retained by the map, the executor side does not need to copy the projected row per record. The class is `private[spark]` and has a single call site. With both sides projecting through identical `UnsafeProjection`s, map lookups compare partition keys by their binary contents, which is consistent with the `InternalRowComparableWrapper` semantics (`RowOrdering`) used to group and de-duplicate `partitionKeys` on the driver. ### Why are the changes needed? The `valueMap` keys and the executor-side lookup keys were `Seq[Any]` compared with Scala `==` element equality, while partition-key grouping/de-duplication uses `InternalRowComparableWrapper` (`RowOrdering`) semantics. The two disagree for `BinaryType`: `Array[Byte]` elements are compared by reference, so every lookup misses and falls back to `nonNegativeMod(hashCode, numPartitions)` -- effectively a random partition per row, since `Array` hash codes are identity-based. As a result, when the non-keyed side of a storage-partitioned join is shuffled into a table partitioned by e.g. `identity(binary_col)`, rows land in partitions that do not match the keyed side, and the join silently drops matches (wrong results, no error). ### Does this PR introduce _any_ user-facing change? Yes, it is a bug fix. Previously, a storage-partitioned join with a `BinaryType` partition key and `spark.sql.sources.v2.bucketing.shuffle.enabled=true` could silently return fewer rows than expected. Now it returns correct results. ### How was this patch tested? A new regression test in `KeyGroupedPartitioningSuite` (`SPARK-59054: shuffle one side: partition keys with binary type`) that joins a v2 table partitioned by `identity` on a binary column with an unpartitioned table. The test was confirmed to fail before the fix (missing join rows) and pass after it. The full `KeyGroupedPartitioningSuite` (105 tests) passes. Note: a NaN-key reproduction was attempted but is not reachable through this path, because `NormalizeFloatingNumbers` wraps floating-point join keys in `KnownFloatingPointNormalized(NormalizeNaNAndZero(...))`, which prevents SPJ from triggering for float/double keys. The `UnsafeRow`-based comparison handles canonical NaN bit patterns correctly regardless. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Fable 5 -- 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]
