apache / apache/datafusion

Speed up SortPreservingMerge by caching per-cursor state across all CursorValues types

Open
#23,855 0 comments 0 reactions 1 assignee Claimed by @zhuqi-lucas View on GitHub
enhancement performance
Dominant language
Rust
Stars
9.3k
Forks
2.4k
Avg merge
3d 7h
Merged PRs (30d)
344

Description

## Is your feature request related to a problem or challenge?

`SortPreservingMerge` (SPM) compares cursor heads in a loser tree — every emitted row triggers `log₂(k)` `T::compare(l, l_idx, r, r_idx)` calls, where `k` is the number of input partitions. Between one advance and the next, a cursor's head does not move, so **the same current-row state gets recomputed on every compare that involves that cursor**.

DataFusion picks a `CursorValues` implementation per sort-key shape:

- Single primitive column (i64, f64, Date, Timestamp, Decimal128, …) → `PrimitiveValues` (already caches).
- Single `Utf8` / `Binary` / `LargeUtf8` column → `ByteArrayValues`.
- Single `Utf8View` / `BinaryView` column → `StringViewArray` (DataFusion's default for strings).
- Multi-column, or types that need arrow's normalized row format → `RowValues` (backed by arrow `Rows`).
- Any of the above with nulls → wrapped in `ArrayValues`.

Today only `PrimitiveValues` uses the `CursorValues::set_offset` hook to cache its current-row value. Every other implementation leaves `set_offset` as the default no-op and re-derives the row bytes from raw buffers on every compare:

- **`ByteArrayValues`**: two `OffsetBuffer::get_unchecked` loads per side per compare, plus a slice construction.
- **`StringViewArray`**: `views().get_unchecked(idx)` load, plus (for external views) a `data_buffers()[buffer_id]` lookup per side per compare.
- **`RowValues`** (multi-column path only): two `rows.row(idx)` offset walks per compare, each chasing through `Arc` → `ArcInner` → `offsets` heap block → `buffer` heap block. Single-column sorts don't hit this path — they use one of the specialized cursors above — but every multi-column sort hits it.
- **`ArrayValues`** (null wrapper wrapping any single-column cursor): two `is_null(idx)` checks per compare unconditionally, even when the batch has no nulls at all.

The chain of pointer/index lookups in the baseline is redundant: within a cursor's stable-head window (many compares between advances), the answer is invariant. Profiling `sort-tpch` shows `SortPreservingMergeStream::poll_next` at ~20% self-time on single-column VARCHAR sort (Q3), and the ratio is similar on multi-column workloads.

## Describe the solution you'd like

Extend the `PrimitiveValues`-style caching pattern to every other `CursorValues` implementation via the existing `CursorValues::set_offset` hook, plus a batch-level fast path in the `ArrayValues` null wrapper.

Two flavors of per-cursor state are cached:

- **Row-level state** (`PrimitiveValues`, `ByteArrayValues`, `StringViewArray`, `RowValues`): refreshed once per `Cursor::advance`, read by `compare` for the current head.
- **Batch-level state** (`ArrayValues`): computed once at cursor construction, read by `compare` as a fast-path branch.

Concretely:

| Cursor type | What compare re-computed each call | What the cache stores |
|---|---|---|
| `ByteArrayValues` | `offsets[idx]`, `offsets[idx+1]` | `(current_start, current_end)` |
| `StringViewArray` → new `StringViewCursorValues` wrapper | view resolution + `data_buffers()[buffer_id]` lookup | `inline_key_fast(view)` when all rows inline; `(current_ptr, current_len)` otherwise |
| `RowValues` | `rows.row(idx).as_ref()` — two offset lookups + slice construction | `(current_ptr, current_len)` into the Arc-owned buffer heap |
| `ArrayValues` | `is_null(l_idx)` + `is_null(r_idx)` on every compare | `has_nulls: bool` — batch-level short-circuit when both sides have no nulls |

Since cached pointers reference Arc-owned heap allocations (not fields of the cursor struct itself), they remain valid across struct moves — no `Pin` needed. `Send`/`Sync` can be hand-implemented for wrappers holding raw pointers, matching the underlying array's guarantees.

This pattern benefits **every** sort workload, not just multi-column:

- Single-column primitive sorts already used `PrimitiveValues`, and get an additional `ArrayValues::has_nulls` fast path.
- Single-column string sorts (`Utf8` / `Utf8View` / `Binary` / `BinaryView`) get the byte-cursor caches.
- Multi-column sorts get the `RowValues` cache — the biggest wins live here because `Rows::row(idx)` chases the deepest pointer chain (Arc → ArcInner → offsets heap → buffer heap) and the merge hot loop hits it more often than any other compare path.

## Describe alternatives you've considered

- **Do nothing.** The current code works and existing sort tests pass. The redundant work is a constant per-compare tax, not a correctness bug.
- **Cache in the merge state machine instead of on the cursor.** Would require the merge to know about every cursor type's storage layout — leaks abstraction and duplicates state across `SortExec` and `SortPreservingMergeExec` (both consume `CursorValues` via the same trait). The `set_offset` hook already exists on the trait as the intended extension point.
- **Precompute all cached state at batch conversion time.** Would allocate O(rows_per_batch) extra memory per batch just for the cached view, and would waste work on rows whose cursor never advances that far because a later batch is loaded early. Refreshing lazily via `set_offset` matches how the merge actually walks the data.

## Additional context

The `set_offset` trait method is already part of `CursorValues`, added when `PrimitiveValues` was optimized. This proposal fills in the other implementations that left it as the default no-op — no new trait methods, no public API change, no allocation change.

PR: #23840

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.