Rich-T-kid opened a new pull request, #11407:
URL: https://github.com/apache/arrow-rs/pull/11407
# Which issue does this PR close?
N/A — standalone perf improvement; no linked issue.
# Rationale for this change
The dictionary comparator returned by `arrow_cmp::make_comparator` was
independent of distinct value count: every slot comparison went through a boxed
value-level comparator, so sorting an N-row `DictionaryArray<_, Utf8>` paid the
full `str::cmp` cost on every one of the `O(N log N)` compares — the dictionary
encoding gave no help on the hot path.
A random-pair benchmark at 1M comparator calls shows the pre-change path
hovering at **~4–6 ms and growing with cardinality** (6.40 ms at 8192), which
is exactly the regime where dictionary encoding is supposed to win.
# What changes are included in this PR?
- `arrow-cmp/src/lib.rs` — rewrote `compare_dict<K>` to precompute per-value
`u32` ranks once at `make_comparator` time:
1. Build three value-level comparators (`left_vs_left`, `right_vs_right`,
`left_vs_right`), all with the same `child_opts(opts)` the original used.
2. Sort each side's value indices independently under its same-side
comparator.
3. Merge-walk the two sorted index lists using the cross-side comparator,
assigning the same `u32` rank to a block of entries that are mutually equal on
either side **and** to blocks that are equal across sides.
4. The returned slot closure becomes `rank_l[left_keys[i] as
usize].cmp(&rank_r[right_keys[j] as usize])` — two indexed loads and a
`u32::cmp` with no boxed dispatch.
- `arrow-cmp/benches/dict_cmp.rs` — new criterion bench that drives
`make_comparator` on `Dictionary<Int32, Utf8>` across cardinalities `[10, 50,
100, 500, 750, 1000, 4000, 8192]`, each with 1M comparator invocations per
iteration.
- `arrow-cmp/Cargo.toml` — adds `criterion` + `rand` dev-deps and the
`[[bench]]` entry.
- Four new differential tests in `arrow-cmp/src/lib.rs` cross-check
`fast(i,j) == naive(i,j)` on every `(i, j)` pair under all four `SortOptions`
combinations, specifically covering:
- duplicate values inside the `values` array on both sides,
- the same key index repeated many times across the `keys` array,
- duplicates mixed with a null value entry,
- the same risks with a primitive (`Int32`) value type.
## Code paths this optimisation touches
- `arrow_cmp::make_comparator` dispatch arm `(Dictionary(_), Dictionary(_))
=> compare_dict::<K>` — the only entry point into this code; dispatched via
`downcast_integer!` on the key type (all `Int{8,16,32,64}` / `UInt{8,16,32,64}`
key types).
- Transitively: anything that calls `make_comparator` on a
`DictionaryArray`. The grep-visible call sites in-tree are:
- `arrow_ord::ord::make_comparator` — re-export; feeds sort/rank/lexsort
for dictionary columns.
- `arrow_ord::sort::lexsort_to_indices` and `rank` — go through the shared
`DynComparator`, so dict columns in lex-sort now get the rank-based path.
- `arrow_select::take` run-end-encoded fast path — unchanged for the REE
case itself, but if the REE values are a `DictionaryArray` this comparator is
what gets built for them.
- No public API surface changes: `make_comparator` signature,
`DynComparator` type, and all four `SortOptions` behaviours are byte-identical
to the previous implementation (verified by the differential tests, and by the
existing `arrow-ord` / `arrow-select` test suites still passing).
## Null handling (unchanged semantics)
- The outer `compare(left, right, opts, …)` wrapper applies
`DictionaryArray::logical_nulls()` which is `key_null OR value_null_at_key`, so
if either the key itself is null or the referenced value is null, the inner
slot closure is never called. That means we only ever index `rank_l`/`rank_r`
with keys pointing at non-null values. Ranks for null values are assigned (via
the same-side sort) but never looked up — they're inert.
- `child_opts(opts)` is applied identically to all three inner comparators,
matching the original implementation's inversion of `nulls_first` under
descending parent order.
# Are there any user-facing changes?
No. This is an internal perf change to `arrow-cmp`. The returned
`DynComparator` is observationally identical to the previous one for every `(i,
j)` pair (asserted exhaustively by the new differential tests), so downstream
`arrow-ord` sort/rank/lexsort behaviour is unchanged.
# Benchmark results
`cargo bench -p arrow-cmp --bench dict_cmp`, `Dictionary<Int32, Utf8>`, 1M
comparator calls per iteration. Reported as `after : before` with the fastest
baseline normalised to 1.00:
| Cardinality | After (ms) | Before (ms) | Ratio (before/after) |
|-------------|-----------:|------------:|---------------------:|
| 10 | 2.01 | 5.06 | 2.52× |
| 50 | 1.95 | 5.42 | 2.78× |
| 100 | 1.96 | 4.21 | 2.15× |
| 500 | 1.90 | 3.95 | 2.08× |
| 750 | 1.92 | 4.04 | 2.10× |
| 1000 | 1.95 | 4.10 | 2.10× |
| 4000 | 1.95 | 5.18 | 2.66× |
| 8192 | 2.17 | 6.40 | 2.95× |
All deltas significant at `p < 0.05`; criterion flags "Performance has
improved" on every size. The rank-based path's cost is now flat across
cardinalities (~2 ms, pure u32 compares) where the old path grew with
cardinality. Rank-table memory is `4 · (L + R)` bytes — ~65 KB at `L = R =
8192`, built once per `make_comparator` call and amortised across every
subsequent slot comparison the caller performs.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
