perf: optimize map_sort for multi-entry string maps (up to 3x faster) - #5901
Conversation
github-actions
Bot
added
enhancement
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Correctness
Summary and compatibility
Reviewed 8282008e388ab05327ab4f7d8dbba94ae034ddab against 4abfd95114d61ad454f9ee269be1615f469f24e2. The prior implementation already handled singleton maps efficiently, but multi-entry rows still sliced their key arrays and invoked Arrow sorting separately. This change reuses string-sort scratch storage for Utf8 keys with Int32 values and bulk-fills offsets for empty visible batches.
The maintained Spark 4.0 implementation sorts keys in ascending order while moving their values with them; its optimizer inserts this normalization for map grouping and repartitioning. Default string ordering is binary UTF-8 ordering. The new comparator agrees with that ordering and reproduces Arrow 59.3's prefix/length comparator and tuple layout. Keeping the tuple layout also preserves Arrow's current unstable permutation for duplicate keys; this does not add a stable-sort guarantee.
The safety and reconstruction boundaries hold: valid MapArray offsets bound every unchecked key access, the four-byte load is guarded by string length, and the nonnegative i32 offsets keep generated indices representable. The global permutation takes whole key/value entries together. Sliced offsets are rebased, map validity and schema metadata are retained, and physical entries below null maps follow the existing sorting/error behavior. Empty batches, all-null batches and the declared-sorted shortcut retain their existing behavior. Nonmatching key/value types continue through Arrow; unsupported nested keys still produce its errors.
The Spark-side type gates continue to reject collated string sorting, and the strict floating-point restriction is unchanged. The maintained Spark 3.5 branch has no MapSort expression; Comet's registration and map-shuffle normalization remain Spark 4.x paths. Maintained 3.4 and 4.1 source branches were unavailable, so the direct source comparison is qualified accordingly.
Validation
All 12 previous unit tests are byte-identical; four additions cover sliced/null/Unicode maps, exact duplicate-key permutations across 48 size/distribution combinations, unsupported keys, and empty visible slices. Rust CI passed all 16 map-sort tests and 1,448 tests overall, with five skipped. Spark 4.1 shuffle CI passed 491 tests, including map keys below an offset, entry-order independence, partition assignment versus Spark, and nested nulls.
Those jobs checked out a8e9f9bc, whose complete source tree equals this PR head. The newer merge preview is distinct; I am not treating the older run as validation of every later base change. At the review cutoff, the head had 53 successful and 10 skipped checks. I performed source and CI verification, without a local native build or benchmark. No verified P1/P2 correctness issue remains in this change.
Performance
The scratch reuse removes per-row key slices and Arrow's temporary index, tuple and output allocations from the selected shape. Scratch is reused across rows and grows with the largest row; the existing batch permutation and single entry take remain. Other shapes pay the added batch type check and fast-path branch while retaining their existing work. Eligible singletons still append directly.
The author reports 122 matched cases against a base that already includes #5887. I checked the inventory: all 18 original slowdown flags have fresh pairs, with reverse-order checks for the two remaining larger fallback flags. Reported 2–10-entry forward normalization is about 3x faster; 2–50-entry fresh pairs improve 33–39%. Initial singleton slowdowns largely disappear in fresh runs. The int-key 16-entry fallback remains +3.50% and +1.99% in the two orders; small effects remain uncertain, and unchanged hash-only controls show appreciable process drift in the original cohort.
The fixtures check exact entries, validity, offsets and schema outside timing. Normalization allocation/drop and combined struct reconstruction are timed; hash-buffer allocation is outside timing and seed reset is inside. Dependency manifests and lockfile are unchanged between the assigned base and head. The timing results and executable hashes are author-reported; I did not independently validate the benchmark binaries or raw samples. This supports the targeted allocation optimization, with the stated microbenchmark limits, rather than an end-to-end Spark speedup claim.
Design
The implementation fits the existing batch-permutation design: one private helper handles the measured string/int shape, while reconstruction and fallback stay shared. Preserving Arrow's comparator avoids introducing a second ordering policy. The empty-visible-batch shortcut still goes through take and rebasing, which is necessary when a slice retains an unused entry prefix. The change leaves hashing, expression registration and feature defaults intact.
Abstraction & complexity
The helper earns its cost by reusing the same scratch allocation across rows. Its explicit tuple and safety comments explain behavior that a simpler index-only unstable sort could change for equal keys. The existing const-generic singleton dispatch remains; there is no additional framework or configuration. The differential Arrow test and independent fixture permutations provide useful coverage for maintaining this specialized code. No additional change is requested before merge.
|
Thanks for reviewing! |
3 participants
Which issue does this PR close?
Closes #5900. Part of #5818; follows the matched normalization/hash benchmarks in #5822 and singleton optimization in #5887.
Rationale for this change
Normalizing multi-entry
map<string,int>batches slices the key array and allocates Arrow sortingindices and prefix tuples for every row. Reuse the prefix-tuple buffer across rows to reduce this
allocation overhead while preserving the exact permutation.
What changes are included in this PR?
Utf8keys withInt32values, reusing(u32 index, u32 prefix, u64 length)scratch storage. Keep Arrow's comparator and tuple layout, including equal-key permutations.
Scratch capacity grows with the largest row. The existing global index append and single
takeremain.map_sortsingleton normalization (18x faster) #5887's singleton eligibility check, batch dispatch and direct index append. Insert themulti-entry path into its const-generic helper; other key/value types retain Arrow validation.
mixed singleton/multi-entry slices, physical entries under null maps, unsupported keys and empty
visible slices. Retain all upstream tests and the existing null/schema/sorted-flag semantics.
How are these changes tested?
-D warningspassed.make format PROFILES=-Pspark-4.0passed, using the existing profile that provides SemanticDB.slowdown flags with CI lower bound above +1% received independent paired measurements, as did
two wide-map wins. Two fallback cases still above +2% received reverse-order pairs.
Matched 2–10-entry forward normalization improved about 66% (3x); matched 2–50-entry cases improved
28–39% in the full cohort and 33–39% in fresh pairs. No larger slowdown persisted across the
confirmations. The int-key 16-entry fallback was +3.50% in a fresh pair and +1.99% in reverse order;
small effects remain uncertain under observed process drift. The int-key 4-entry case fell from
+6.09% to +0.11% in reverse order. Initial singleton +8–9% flags did not reproduce: fresh matched
singleton pairs were -0.08% and +1.72%. These are microbenchmarks, not end-to-end Spark measurements.
Measurement cohort and limits
Base:
4abfd95114d61ad454f9ee269be1615f469f24e2, including merged singleton PR #5887 (ca223a8d4).All 122 baseline measurements completed before the integrated production source was applied.
Both cohorts used byte-identical benchmark files, dependency manifests and lockfile.
Apple M4 Max, 64 GB, Rust 1.97.1, release/thin LTO/codegen-units=1, no RUSTFLAGS override.
Full cohorts: 50 samples, 1 s warmup, 2 s target measurement (Criterion extends long cases).
Independent pairs: 100 samples, 3 s warmup, 5 s target measurement. No heavy build or benchmark ran
concurrently. Narrow within-process CIs do not capture desktop/background and cross-process drift.
All unfavorable original measurements are retained alongside confirmations.
Inputs and correctness checks are untimed. Normalization output allocation/drop and combined
struct reconstruction/drop are timed. Hash-buffer allocation is untimed and seed reset is timed.
Hash-only forward/reversed cases use the same normalized input as repeated controls. Each batch
has 8,192 visible rows; physical entries under null maps count toward entry rates. No phase subtraction
is used. Hashing implementation and nested-hashing defaults are unchanged.
Complete integrated comparison (122 cases)
Time is ms per 8,192-row batch, with 95% CI; unsupported-key and buffer-reset cases are per call.
The change column is Criterion's cross-process mean estimate; confirmations follow below.
Fresh independent pairs and reverse-order confirmations
Benchmark executable identities (SHA-256)
map_sort:d93a75c7df21c0b57fb994664c6024b8aa07bea85939f08d7b0e6279689c6e8fhash:05739bd6b3ef97c63ed814943cd994e205212abb9a95c8484c977ebdc77e3207map_sort:47a0a73bbedfd8c6339c0e3626d1d0af6ffb9e6fdee679c8c26896905147f7a3hash:ccc68e4d95296c241e222a7c5a552c8a57c30b5365fccbe72ff097a5400e7dd3