Skip to content

fix: rebase map offsets in mapsort so sliced maps do not overrun entries - #5630

Open
viirya wants to merge 1 commit into
apache:mainfrom
viirya:fix-map-sort-sliced-offsets
Open

fix: rebase map offsets in mapsort so sliced maps do not overrun entries#5630
viirya wants to merge 1 commit into
apache:mainfrom
viirya:fix-map-sort-sliced-offsets

Conversation

@viirya

@viirya viirya commented Sep 2, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Closes #5629.

Rationale for this change

spark_map_sort builds its take indices only for the maps that are visible, so
the sorted entries it produces are indexed from zero. It then passed the input
offsets straight to MapArray::try_new:

let sorted_entries = take(maps_arg_entries, &indices, None)?;
...
MapArray::try_new(
    Arc::clone(map_field),
    maps_arg.offsets().clone(),   // original offsets, not rebased
    sorted_map_struct.clone(),
    ...
)

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 (DataFusion's limit uses
batch.slice(skip, ...)), and Spark 4.0+ inserts MapSort on its own, so this fails
a real query today — a group-by on a map column below an OFFSET, with no explicit
mapsort call and no configuration changes.

What changes are included in this PR?

  • map_sort.rs: rebuild the output offsets from the per-map lengths rather than
    reusing the input offsets. Empty maps still contribute an offset entry, so a map
    array containing empty maps keeps the right offset count (the previous continue
    skipped the whole iteration).

How are these changes tested?

  • New Rust unit test test_sliced_map_offsets_are_rebased: slices a two-map array to
    drop 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_sort passes 9 tests.
  • New end-to-end test in CometMapExpressionSuite: a group-by on a map column below
    LIMIT ... OFFSET ..., compared against Spark. CometMapExpressionSuite passes 23
    tests.

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 InsertMapSortInRepartitionExpressions as a
second 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 add
the repartition-on-map coverage to #5567, where that path becomes reachable.

`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 sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

mapsort on a sliced map fails with "Max offset exceeds length of entries"

2 participants