fix: rebase map offsets in mapsort so sliced maps do not overrun entries - #5630
Conversation
sunchao
left a comment
There was a problem hiding this comment.
Reviewed 8eddd6d6a1172a09c112f348557342cd088171aa against the reported base 8729f6e6adf7091e18a48670e790d4ba8fd41e51. No actionable P1/P2 findings. Rebuilding one output boundary per visible map fixes the sliced-offset overrun while preserving empty maps, validity, and key/value pairing.
A focused Arrow 58.4.0 probe using the unchanged production function bodies passed 3 tests, including all 36 slice windows of mixed nonempty/empty/null maps and nullable values. The base function reproduces the [2, 4] overrun; this head returns the correctly rebased result. This is array-level validation with substituted DataFusion wrappers, not a full Comet or Spark execution. I did not run the supplied Rust/Scala suites or benchmarks.
Both Linux and macOS lint failures are cargo fmt --check differences in the new Rust test. Please format those before merging; other CI checks are still running. The Scala regression would also benefit from native OFFSET/MapSort plan assertions, since checkSparkAnswer alone permits fallback.
`spark_map_sort` builds its `take` indices only for the maps that are visible, so the sorted entries it produces are indexed from zero, but it passed the input offsets straight to `MapArray::try_new`. A sliced `MapArray` keeps its original entry offsets, so for two two-entry maps sliced to drop the first, the offsets are `[2, 4]` against 2 taken entries and Arrow rejects the result with "Max offset of 4 exceeds length of entries 2". A native OFFSET is enough to produce that slice, and Spark 4.0+ inserts `MapSort` on its own, so a group-by on a map column below an OFFSET fails the query today. Rebuilds the output offsets from the per-map lengths instead. Empty maps still contribute an offset entry, so a map array containing empty maps keeps the right offset count. Co-authored-by: Claude Code <noreply@anthropic.com>
8eddd6d to
983f6fd
Compare
sunchao
left a comment
There was a problem hiding this comment.
Thanks for the formatting update. Checked 983f6fd5: the production code is unchanged, both lint jobs now pass, and I found no new issue in the base-integration check. I reused the earlier Arrow-only probe evidence without rerunning tests. No Spark/JNI execution is claimed; the remaining CI is still pending.
|
Thanks @sunchao ! |
`CometShuffleExchangeExec.supportedHashPartitioningDataType` rejected struct, array and map partitioning keys, so any query repartitioning on a nested column fell back to Spark for the whole shuffle. The comment said "Native code does not support hashing complex types, see hash_funcs/utils.rs", but that file hashes nested types recursively (struct fields, list elements, map keys and values), and shuffle partitioning shares that kernel, and Spark's seed, with the `hash` expression. Adds the recursive struct/array/map cases to the gate, behind `spark.comet.shuffle.native.partitioning.hash.nested.enabled` (default true). Nesting is checked recursively through the same predicate, so a leaf type that cannot be hashed natively disqualifies the whole key and the shuffle still falls back: - collated strings, which Comet hashes as raw bytes (see apache#1947 / apache#4035, where rows equal under the collation reached different partitions and a downstream collation-aware DISTINCT produced a wrong answer) - CalendarInterval, which the native hasher has no branch for (apache#5059) Map keys are additionally restricted to Spark 4.0+. Map entry order is not semantically meaningful, so two equal maps must hash alike, and Spark 4.0+ normalizes a map shuffle key by wrapping it in `mapsort(...)`. Earlier versions insert no such normalization, so Comet would hash physical entry order. When the `mapsort` itself is not convertible -- CometMapSort supports scalar map keys only -- the existing expression check fails and the shuffle falls back. The config defaults to false. The native hasher only vectorizes nested shapes whose leaves are primitives; `array<struct<...>>` and a map inside a struct fall through to a per-element path that re-enters `create_murmur3_hashes` for every element, so enabling this by default before measuring could make these shuffles slower than letting Spark do them. `CometFuzzTestSuite`'s "distribute by single column (complex types)" keeps its existing expectation, since the keys still fall back by default, and additionally asserts that they are admitted with the config enabled. Also corrects the comment in `CometFuzzTestBase` claiming that file has no nested complex types -- `generateSchema` does add `struct<array<..>>` and `array<struct<..>>` when both array and struct generation are on; maps are the only thing missing. Also covers the repartition-on-map route into the sliced-map `mapsort` defect fixed in apache#5630: a native OFFSET below a map shuffle key, asserting that both the exchange and the offset stay native so the sliced map actually reaches the native mapsort. Co-authored-by: Claude Code <noreply@anthropic.com>
Which issue does this PR close?
Closes #5629.
Rationale for this change
spark_map_sortbuilds itstakeindices only for the maps that are visible, sothe sorted entries it produces are indexed from zero. It then passed the input
offsets straight to
MapArray::try_new:A sliced
MapArraykeeps its original entry offsets, so for two two-entry mapssliced to drop the first, the offsets are
[2, 4]against 2 taken entries and Arrowrejects the result with
Max offset of 4 exceeds length of entries 2.A native OFFSET is enough to produce that slice (DataFusion's limit uses
batch.slice(skip, ...)), and Spark 4.0+ insertsMapSorton its own, so this failsa real query today — a group-by on a map column below an OFFSET, with no explicit
mapsortcall and no configuration changes.What changes are included in this PR?
map_sort.rs: rebuild the output offsets from the per-map lengths rather thanreusing the input offsets. Empty maps still contribute an offset entry, so a map
array containing empty maps keeps the right offset count (the previous
continueskipped the whole iteration).
How are these changes tested?
test_sliced_map_offsets_are_rebased: slices a two-map array todrop the first map, asserts the slice really does keep the original offsets, and
checks the sorted keys and values.
cargo test -p datafusion-comet-spark-expr map_sortpasses 9 tests.CometMapExpressionSuite: a group-by on a map column belowLIMIT ... OFFSET ..., compared against Spark.CometMapExpressionSuitepasses 23tests.
I checked that both tests actually guard the fix by reverting the offset rebase and
rebuilding: the Rust test and the end-to-end test both fail, with the Arrow error
above.
Additional context
Found by @sunchao while reviewing #5567, which admits nested types as native shuffle
hash partitioning keys and would add
InsertMapSortInRepartitionExpressionsas asecond route into this code. This fix is independent of that PR — the group-by path
reproduces on current
main— so it is split out here. Once this merges I will addthe repartition-on-map coverage to #5567, where that path becomes reachable.