perf: optimize map_sort singleton normalization (18x faster) - #5887
Conversation
github-actions
Bot
added
enhancement
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Correctness
Reviewed cf22afe92e5d504adf72917a3a12473f9fb3c578 against de1eb4f86c12af0895784c93e1f152c705f6ef0e. I found no P1/P2 issue.
Previously, every nonempty map created a key slice and dispatched Arrow sorting, even when it contained one entry. This change appends that entry's original index directly for supported flat key types. Wider maps still use the same Arrow sort and the same global permutation. The final take, rebased offsets, map validity buffer, field metadata and sorted flag remain intact. Entries physically stored beneath null maps are still processed unless the existing all-null early return applies. Empty batches and already-sorted inputs keep their existing early returns.
On the maintained Spark 4.0 branch, MapSort orders key/value pairs together and preserves the map type in interpreted and generated execution. A singleton has no ordering decision, so retaining its original pair preserves that behavior without altering NaN payloads or signed zero. The locked Arrow 59.3 sources confirm that the eligible flat types use infallible sort implementations. Lists, dictionaries, structs and other excluded types continue through Arrow, including its singleton validation errors. The existing Comet restrictions for nested keys, collated strings and strict floating-point behavior are unchanged. This does not expand Spark compatibility. Spark 3.5 has no corresponding MapSort insertion and Comet still excludes map shuffle keys there. Maintained Spark 3.4/4.1 source branches were unavailable.
Validation
The three new Rust tests supplement nine existing map-sort tests. They check explicit permutations for mixed and singleton slices, nonzero offsets, hidden entries under null maps, nullable values, schema metadata, shared validity buffers, binary types, NaN bits, signed zero, unsupported nested-key errors and early returns. These are direct kernel tests rather than constant-folded SQL expressions.
At the September 12, 22:45 UTC refresh, CI had 56 successful and 10 skipped checks. The Rust job passed all 12 map-sort tests, all 11 Murmur3 tests and 1,418 tests overall. The Spark 4.1 expression job passed 1,384 tests, including the existing sliced-map regression. That integration test checks Spark answers without independently asserting native MapSort execution. Both jobs checked out merge 8e64630b7a87cdb83f838d5afd2c435fc78f9ef4, whose parents are the assigned base and head. All five changed files and inspected supporting paths match the head. Five inherited base files make the complete trees different. I did not run a local product build or benchmark.
Performance
The optimization removes a key-array slice, sort dispatch and local index allocation for each eligible singleton. It retains the global permutation and output take, so this is not a zero-copy result. A preliminary offset scan adds linear work to eligible batches without singletons. The const-specialized fallback avoids adding a singleton branch to their inner sorting loop. Mixed batches stop the preliminary scan at the first singleton.
The shared harness separates normalization, hashing and combined execution. Its added cases include sliced inputs, mixed cardinalities, null maps and values, long Unicode strings, empty maps and unsupported singleton keys. Input construction and correctness checks are untimed. Hash-buffer resets and normalization allocation/drop remain timed. The unsupported-key case measures first-error cost, not batch throughput.
The author reports 19–22x faster matched singleton normalization, 18.4x in an independent confirmation and 13.84x for one combined singleton case. The published tables retain the initial regression flags and subsequent independent/reversed-order measurements. These are author-reported microbenchmarks from a DataFusion 55.0 development cohort. The reviewed head and CI instead lock DataFusion 55.1. I have not independently reproduced those timings or established an exact-head speedup, and they do not measure end-to-end Spark query performance. The audit documentation identifies the development cohort explicitly.
Design
The change optimizes the existing normalization boundary without changing hashing or shuffle defaults. Keeping one global permutation preserves the established behavior for slices and null-map storage. Restricting the shortcut to flat types is justified by the fallible child-ranking paths for nested types. An unconditional singleton shortcut would suppress existing errors. The targeted eligibility check and shared output construction keep the behavioral change small.
Abstraction & complexity
The private const-generic helper gives both paths one loop implementation while removing the extra branch from the fallback specialization. It introduces no new configuration or public API. Benchmark additions reuse the existing shared fixtures and stage definitions, with independent expected permutations outside timing. I found no actionable abstraction or complexity issue.
andygrove
left a comment
andygrove
left a comment
There was a problem hiding this comment.
The singleton fast path looks right to me and the 19-22x on normalize_only/map/1 is well clear of any plausible harness noise. I checked the key type allowlist against arrow 59.3 and it lines up exactly with the infallible arms of sort_to_indices, so the correctness argument for skipping the sort holds.
|
Thanks @sunchao @andygrove |
3 participants
Which issue does this PR close?
Closes #5886. Follow-up to #5822.
Rationale for this change
A singleton map already has its only possible ordering, but normalization still creates a key slice and invokes Arrow sorting for every row. Avoiding that work speeds up matched singleton normalization by 19–22x in the full measurement and 18.4x in an independent forward-order confirmation; normalization plus hashing improves by 13.84x in the full forward singleton case.
What changes are included in this PR?
Skip per-row sort dispatch only for supported flat singleton key types. A batch check chooses a const-specialized loop so wide maps retain a fallback without the added per-row singleton branch. Keep nested and unsupported types on Arrow's validation path, including singleton errors. Preserve global permutation/take, rebased offsets, physical entries beneath null maps, schema metadata, and existing early returns. Hashing implementation and nested-hashing defaults are unchanged.
Extend the shared matched benchmark harness with null, sliced, mixed-cardinality, long Unicode, empty, singleton, and unsupported-key cases. Record the local performance audit.
How are these changes tested?
The unchanged nine map-sort tests and three new tests pass (12 total), covering manually expected permutations, exact values, schema metadata, null buffers, nonzero offsets, nested keys, unsupported errors, NaN payloads, signed zero, and binary keys. All 11 Murmur3 tests pass. Both benchmark smoke runs pass (92 cases). Full
makesucceeds (native and JVM packaging; JVM tests skipped by the build target). Workspace all-target Clippy passes with warnings denied, and Rust formatting passes.Benchmark environment: Apple M4 Max, macOS 26.6.2, Rust 1.97.1, thin LTO, one codegen unit, DataFusion 55.0.0. Both sides use commit 3810936 plus identical benchmark additions and the same dependency lock (SHA-256
02a7248aa329e1f98c7cce0ae921601ca66cefae66dd2c08631c98d3ae8337cb); only the optimized side includes the production change. These are microbenchmarks, not end-to-end Spark query speedups.Criterion: 8192 visible rows, seed 42, 100 samples, 3-second warmup, 5-second requested measurement, 95% confidence, sequential execution with no concurrent builds. Input construction and checks and hash-buffer allocation are untimed; seed reset, normalization allocation/drop, and combined struct construction/drop are timed. No differences between stage or process timings are presented as exact cost attribution. The unsupported-key benchmark fails at the first row and is not batch throughput.
All 92 cases were measured. Fourteen initial regression flags were independently remeasured. The three residual approximately 1.4% flags were measured again with execution order reversed: hash-only directions reversed and the combined case returned within the 1% noise threshold. No stable meaningful regression was observed; confidence intervals within a process do not cover machine drift. All initial flags remain visible below.
All 92 cases: timings and 95% confidence intervals
Timings use Criterion's slope estimate when available, otherwise its mean. Speedup is the ratio of those point estimates. Regression confirmation separately uses Criterion's relative mean change, so its percentages need not equal the slope ratio.
Independent and reversed-order confirmation
Positive means optimized is slower. Reverse-order estimates are reciprocally converted to the same direction.
Criterion timing and change output for all 92 cases