peter-toth commented on code in PR #58345:
URL: https://github.com/apache/spark/pull/58345#discussion_r3873608829
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala:
##########
@@ -407,10 +406,15 @@ object ShuffleExchangeExec {
assert(k.isGrouped,
s"Expected a grouped KeyedPartitioning on ${k.expressions}, but got
${k.numPartitions} " +
"partition keys with duplicates among them")
+ // Project the partition keys to UnsafeRows so that map lookups
compare them by value
+ // (e.g. binary keys by content, NaNs as equal), consistently with
both the
+ // InternalRowComparableWrapper semantics used to group
`partitionKeys` and the lookup
+ // keys produced by `getPartitionKeyExtractor` below.
+ val projection = UnsafeProjection.create(k.expressionDataTypes.toArray)
Review Comment:
**Finding 1.** `UnsafeRow` byte equality is strictly finer than the equality
that decided which keys share a partition, so a lookup can still miss the map —
and for one shape it now misses where base hit.
Byte-equal implies `RowOrdering`-equal, but not the reverse. `partitionKeys`
is grouped and de-duplicated with `InternalRowComparableWrapper`, i.e.
`RowOrdering`. So any key pair the driver collapsed but this projection splits
reaches `valueMap` under only one of its two byte forms, and a row carrying the
other form falls to `Utils.nonNegativeMod(key.hashCode, numPartitions)` and
lands in a partition that cannot hold its match. That is the same lost-match
failure this PR fixes.
Three pairs where `RowOrdering` says equal and the projection says different
— measured, same declared type, same projection:
| pair | `RowOrdering` | `UnsafeRow` | base `Seq[Any]` `==` |
|---|---|---|---|
| `-0.0` vs `0.0` | equal | **different** | **equal** |
| `Double.NaN` vs `longBitsToDouble(0x7ff8000000000001L)` | equal |
different | different |
| `'aa'` vs `'AA'` at `UTF8_LCASE` | equal | different | different |
The first row is a regression: base matched those, this does not.
`UnsafeRow` canonicalises neither `-0.0` nor NaN — that is
`NormalizeFloatingNumbers`' job, and `UnsafeWriter.writeDouble` is a bare
`Platform.putDouble` — and its bytes are never collation-aware. So the comment
above claims more than the code gives: the lookup agrees with
`InternalRowComparableWrapper` on binary content, not on `RowOrdering`.
**Reachability.** An ordinary float, double or collated *join* key cannot
get here. `NormalizeFloatingNumbers` wraps it in
`KnownFloatingPointNormalized(NormalizeNaNAndZero(...))`, and a collated
equi-join key becomes `collationkey(...)`; either way
`KeyedShuffleSpec.keyPositions` finds nothing for
`e.references.head.canonicalized` in `distribution.clustering`, the bitset is
empty and SPJ never triggers. I checked both: `identity` on a `DoubleType`
column plans 2 shuffles, and so does `identity` on a `StringType("UTF8_LCASE")`
column. But `keyPositions` matches the transform's *argument*, not its result
type — so a transform whose `resultType()` is `DoubleType` over a plain
`LongType` join key goes straight through, and `k.expressionDataTypes` is then
`Seq(DoubleType)`.
**Measured.** A `ScalarFunction` with `inputTypes() = Array(LongType)` and
`resultType() = DoubleType`, mapping id 1 to `-0.0`, id 2 to `0.0` and id 3 to
`3.0`; `items` holding ids 1..3 partitioned by it, `purchases` unpartitioned
with the same ids, joined on `id`:
| arm | shuffles | rows |
|---|---|---|
| base `c3d9631787e` | 1 | 3 |
| this PR `11f8171815f` | 1 | **2** — id 2 lost |
| this PR, `v2.bucketing.shuffle.enabled=false` | 2 | 3 |
**Fix.** Let the map key be the producer's own wrapper, so the lookup and
the grouping share one equivalence by construction. Driver:
```scala
val valueMap = k.partitionKeys.zipWithIndex.toMap[Any, Int]
new KeyGroupedPartitioner(valueMap, k.numPartitions)
```
Executor, in place of the projection:
```scala
val wrap = InternalRowComparableWrapper
.getInternalRowComparableWrapperFactory(k.expressionDataTypes)
// ... fill partitionKeyRow as you do now, then:
wrap(partitionKeyRow)
```
That needs `InternalRowComparableWrapper` to be `Serializable`, with
`structType` and `ordering` `@transient` and re-derived from the shared caches
on first use after deserialization — the generated ordering cannot cross the
wire. With those three hunks the suite is 105/105 and the probe above returns 3
rows. `Partitioner.scala`'s new contract also gets shorter, because the map
keys simply *are* `partitionKeys`.
The trade-off is per-row cost: `InternalRowComparableWrapperBenchmark`
reports ~98 ns/row for `toSet` over 200k wrappers, against a single Murmur pass
for `UnsafeRow.hashCode`. If you would rather keep `UnsafeRow` and take the
`-0.0` narrowing knowingly, then this comment needs to stop claiming
consistency with `InternalRowComparableWrapper` and name what is excluded
instead.
Either way, a regression test wants a connector function whose
`resultType()` differs from its `inputTypes()`. `InMemoryBaseTable` needs a
whitelist entry and a `getKey` case for that — both are closed matches, at
`InMemoryBaseTable.scala:210` and `:319`.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -1691,6 +1691,43 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase with
}
}
+ test("SPARK-59054: shuffle one side: partition keys with binary type") {
+ val items_partitions = Array(identity("id"))
+ createTable(items, Array(
+ Column.create("id", BinaryType),
+ Column.create("name", StringType),
+ Column.create("price", DoubleType)), items_partitions)
+
+ sql(s"INSERT INTO testcat.ns.$items VALUES " +
+ "(X'0101', 'aa', 40.0), " +
+ "(X'0202', 'bb', 10.0), " +
+ "(X'0303', 'cc', 15.5), " +
+ "(X'0404', 'dd', 20.0)")
+
+ createTable(purchases, Array(
+ Column.create("item_id", BinaryType),
+ Column.create("price", DoubleType)), Array.empty)
+ sql(s"INSERT INTO testcat.ns.$purchases VALUES " +
+ "(X'0101', 42.0), (X'0101', 44.0), (X'0202', 11.0), (X'0202', 19.5), " +
+ "(X'0303', 26.0), (X'0303', 30.0), (X'0404', 50.0), (X'0404', 60.0)")
+
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
Review Comment:
**Finding 5.** The four sibling tests just above (`SPARK-41471: shuffle one
side: ...`) loop both conf states and assert 2 shuffles in the off arm. Worth
doing here too — with the feature off both sides shuffle by hash on the binary
column, so the same 8 rows must come back, which pins the answer as a property
of the query rather than of the partitioner. I ran this shape and it passes on
the PR head:
```scala
Seq(true, false).foreach { shuffle =>
withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key ->
shuffle.toString) {
val df = createJoinTestDF(Seq("id" -> "item_id"))
val shuffles = collectShuffles(df.queryExecution.executedPlan)
if (shuffle) {
assert(shuffles.size == 1, "only shuffle one side not report
partitioning")
} else {
assert(shuffles.size == 2, "should add two side shuffle when
bucketing shuffle one side" +
" is not enabled")
}
checkAnswer(df, Seq(
Row(Array[Byte](1, 1), "aa", 40.0, 42.0),
// ... unchanged
Row(Array[Byte](4, 4), "dd", 20.0, 60.0)))
}
}
```
--
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]
