Skip to content

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

Merged
viirya merged 1 commit into
apache:mainfrom
viirya:fix-map-sort-sliced-offsets
Sep 2, 2026
Merged

fix: rebase map offsets in mapsort so sliced maps do not overrun entries#5630
viirya merged 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.

@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.

`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>
@viirya
viirya force-pushed the fix-map-sort-sliced-offsets branch from 8eddd6d to 983f6fd Compare September 2, 2026 16:11

@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.

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.

@viirya
viirya merged commit 88176ea into apache:main Sep 2, 2026
73 checks passed
@viirya
viirya deleted the fix-map-sort-sliced-offsets branch September 2, 2026 18:30
@viirya

viirya commented Sep 2, 2026

Copy link
Copy Markdown
Member Author

Thanks @sunchao !

viirya added a commit to viirya/arrow-datafusion-comet that referenced this pull request Sep 2, 2026
`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>
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