viirya opened a new pull request, #5567:
URL: https://github.com/apache/datafusion-comet/pull/5567
## Which issue does this PR close?
Closes #5566.
## Rationale for this change
`CometShuffleExchangeExec.supportedHashPartitioningDataType` rejected struct,
array and map partitioning keys, so a query repartitioning on a nested
column fell
back to Spark for the whole shuffle. The comment justified this with:
> Native code does not support hashing complex types, see hash_funcs/utils.rs
but that file hashes nested types recursively (struct fields,
`List`/`LargeList`/`FixedSizeList` elements, map keys and values), and
shuffle
partitioning calls the same `create_murmur3_hashes` entry point, with the
same
seed, that the `hash` expression uses -- where nested types are already
allowed
and covered by `CometHashExpressionSuite`.
## What changes are included in this PR?
- `CometShuffleExchangeExec`: add recursive `StructType` / `ArrayType` /
`MapType`
cases to the hash partitioning gate, and replace the stale comment.
Recursion
goes through the same predicate, so a leaf type that cannot be hashed
natively
disqualifies the whole key and the shuffle still falls back -- this covers
collated strings (#1947 / #4035) and `CalendarInterval` (#5059) without
needing
separate checks.
- Map keys are admitted only on Spark 4.0+. Map entry order is not
semantically
meaningful, so equal maps must hash alike; Spark 4.0+ normalizes a map
shuffle
key with `mapsort(...)` (#1941) while earlier versions do not, and Comet
would
otherwise hash physical entry order. When the `mapsort` is not convertible
(`CometMapSort` supports scalar map keys only) the existing
partitioning-expression check already forces a fallback.
- New config `spark.comet.shuffle.native.partitioning.hash.nested.enabled`
(default `true`) to restore the old behavior.
## How are these changes tested?
`CometNativeShuffleSuite`, run on two profiles:
- Spark 4.1 (Scala 2.13): 46 tests, 46 succeeded
- Spark 3.5 (Scala 2.12): 43 succeeded, 3 canceled by `assume(isSpark40Plus,
...)`
(the map-key and collation cases)
Also green: `CometShuffleSuite` (44), `DisableAQECometShuffleSuite` (43),
`CometHashExpressionSuite` (37).
New tests cover struct, array, two-level (`struct<array<..>>` and
`array<struct<..>>`), and a four-level key mixing all three recursive
branches:
```
struct< a: array< struct< m: map<string, array<int>>, s: string > >, i: int >
```
including empty-array and null-key rows. Fallback tests cover a map key on
Spark
3.x, a map whose own key is nested, a collated string nested inside a
struct, and
an interval leaf inside a struct. The existing `native shuffle on nested
array`
test previously asserted that a nested key falls back; it now asserts native
by
default and fallback when the config is disabled.
Two notes on how this was verified beyond the assertions:
- `checkShuffleAnswer` asserts `shuffleType == CometNativeShuffle` at the
plan
level. To confirm the nested key really is hashed in native code, I also
inspected the executed plan and the exchange metrics for a deep key, which
report native shuffle-writer counters (`repart_time`, `encode_time`,
`interleave_time`, `input_batches`) that the Spark and columnar paths do
not
produce.
- The partition-assignment test compares `spark_partition_id()` per row
against
Spark for shallow and deep keys. I sanity-checked that it can fail by
perturbing the native struct hash, which makes it report every partition id
shifted by one.
--
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]
