fix: rebase map offsets in mapsort so sliced maps do not overrun entries - #5630
fix: rebase map offsets in mapsort so sliced maps do not overrun entries#5630viirya wants to merge 1 commit into
Conversation
`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>
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.
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.