peter-toth opened a new pull request, #58820: URL: https://github.com/apache/spark/pull/58820
### What changes were proposed in this pull request? Backport of [SPARK-59054](https://issues.apache.org/jira/browse/SPARK-59054) (apache#58345) to `branch-4.0`, cherry-picked from `branch-4.2` [`668683054db`](https://github.com/apache/spark/commit/668683054db5e634c8d9ccc664f62b2d1c769b30), the lowest branch that has it. The `branch-4.1` backport is apache#58819. The fix makes the partition-key lookup in `KeyGroupedPartitioner` use the same equivalence that grouped the partition keys: - In `ShuffleExchangeExec.prepareShuffleDependency`, the driver-side `valueMap` is keyed by `InternalRowComparableWrapper`s over the partitioning's expression data types, instead of `Seq[Any]` from `InternalRow.toSeq`. - The executor-side lookup key in `getPartitionKeyExtractor` is built by evaluating the bound partition expressions into a reused `GenericInternalRow` and wrapping it the same way. - `InternalRowComparableWrapper` becomes `Serializable`. `structType` and `ordering` cannot cross the wire (the ordering may be generated code), so they are transient and re-derived from the shared caches on first use after deserialization. This is what allows the wrappers to be shipped to executors inside the partitioner. - `KeyGroupedPartitioner` takes `Map[Any, Int]` and looks keys up with `getOrElse`, removing the `ArraySeq` normalization and the `getOrElseUpdate` map mutation. **Two things differ from the 4.2 commit, both for the same reason.** `getInternalRowComparableWrapperFactory` was added by SPARK-54383, which is not on this branch, and I did not want to drag an optimization onto a maintenance branch: - `InternalRowComparableWrapper` keeps its public two-argument constructor rather than gaining a private four-argument one plus a deprecated delegate. `_structType` and `_ordering` are transient `var`s filled on first use, which is the same behaviour the 4.2 class has after deserialization. - Both call sites construct wrappers directly instead of through a hoisted factory. On the executor side that is one construction per row, but the fix still makes that path cheaper than the branch has today, because it hoists `bindReferences` out of the per-row closure, which the pre-fix code called for every record. Three smaller adaptations: - The incoming `getPartitionKeyExtractor` hunk carries a `ShufflePartitionIdPassThrough` arm. That class does not exist on this branch, so the arm is dropped. - `InMemoryBaseTable` keeps this branch's `Literal` name for the connector literal, which 4.2 imports as `V2Literal`. - `withFunction` does not exist in `KeyGroupedPartitioningSuite` here, so `UnboundSignedZerosFunction` is registered in the suite's own `functions` list and the wrapper block unwrapped. Everything else is the upstream commit. `Partitioner.scala` and `transformFunctions.scala` applied clean and were verified identical to upstream with a diff of diffs, apart from one scaladoc reference the three-way merge adjusted from `KeyedPartitioning` to this branch's `KeyGroupedPartitioning`. ### 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 and de-duplication use `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. So 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. **Confirmed on this branch.** With the fix reverted and the tests kept, `SPARK-59054: shuffle one side: partition keys with binary type` fails with missing join rows. The other two new tests pass either way on `branch-4.0`: struct field names cannot disagree here because `InternalRowComparableWrapper.mergePartitions` wraps both sides with one expression list, and `-0.0 == 0.0` already held under Scala `Double` equality. They are kept as upstream wrote them. ### 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? The three tests that ship with the upstream commit, run on this branch. `KeyGroupedPartitioningSuite` is green at 88 tests, and `dev/lint-scala` passes. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Opus 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]
