Skip to content

[Enhancement]: Speed up storage.MergeSort with a proper k-way merge #51981

Description

@bigsheeper

Is there an existing issue for this?

  • I have searched the existing issues

What would you like to be added?

Three changes to storage.MergeSort (internal/storage/sort.go:326-490), the k-way merge used by mix compaction:

  1. Keep only k entries in the heap. enqueueAll currently pushes every qualifying row of the current record into the priority queue, so the heap holds the number of in-flight rows rather than the number of readers. Since each input record is already sorted by the merge key, the standard k-way form (one entry per reader, pop one and push one) applies. At 30 readers x 4096 rows/record this takes the heap from ~246k entries down to 30 — log2 from ~18 to ~5, and small enough to stay in L1.

  2. Stop re-resolving the key column on every comparison. The comparator does recs[x.ri].Column(fid).(*array.Int64).Value(x.i) per side per comparison — a map lookup plus a type assert whose result is invariant for a given record. Pre-extract the key column once per record in advanceRecord (Int64Values() for int64, the *array.String pointer for varchar; both O(1) and copy-free).

  3. Remove the per-row heap allocation. pq.Enqueue(&index{ri: ri, i: j}) escapes once per row, and PriorityQueue[T] stores []*T while container/heap's Push(x any) boxes regardless. A small typed value heap over the existing rowIndex{ri, i int32} (added by enhance: speed up sort compaction in storage.Sort #50817) removes it.

storage.PriorityQueue then has no remaining users repo-wide and is removed, along with the endPositions bookkeeping.

The sibling function Sort in the same file already received (2) and (3) in #50817; MergeSort was not touched.

Why is this needed?

MergeSort is the default mix-compaction path: dataNode.compaction.useMergeSort defaults to true and maxSegmentMergeSort to 30, so mix compaction over already-sorted segments goes through it (mix_compactor.go:472-495 -> merge_sort.go:135).

Measured on a standalone reproduction of the merge machinery outside the repo (records shaped like simpleArrowRecord, 8 int64 columns, ascending PK per record), median of 3 runs, -cpu 1, i7-8700:

shape current this change
8 readers x 4 records x 4096 rows (131k rows) 802 ns/row 49 ns/row (16.2x)
30 readers x 2 records x 4096 rows (246k rows) 1059 ns/row 80 ns/row (13.2x)

Allocations at 246k rows: 245,818 -> 28.

Scope of that number: it isolates the merge machinery. Adding a representative per-row copy of 8 int64 columns, standing in for the untouched rb.Append/appendValueAt, gives 1307 -> 233 ns/row (5.6x) at 30 readers. On a collection with wide vector columns appendValueAt copies kilobytes per row, so the end-to-end compaction gain will be lower; I have not measured that shape and am not claiming it.

Anything else?

Behavior preservation. The comparator is a strict total order (key, then ri, then i) and every heap element is a distinct (ri, i), so the pop sequence is determined by the comparator alone, independent of the heap implementation. I verified the full emitted (ri, i) sequence is identical between the current and proposed forms across 6 shapes x 5 seeds, plus a case with duplicate keys spanning readers and record boundaries.

One deliberate behavior change. A proper k-way merge relies on each input record being sorted by the merge key. That reliance is not new — the existing endPositions mechanism already depends on it: if a record were unsorted, the row with the smallest key could be dequeued first, triggering advanceRecord while stale entries from the now-invalid record remain in the queue (RecordReader.Next is borrow-scoped per record_reader.go:20-21). What differs is the failure mode: today that situation panics with index out of range, whereas a bare k-way merge would silently emit rows out of order. To avoid trading a loud failure for silent data corruption, this change adds an O(1)-per-row check that the popped key never decreases, returning a merr error instead.

This is not a fix for #48322 / #48449. The root cause there is that sortMergeAppicable (mix_compactor.go:474-486) only checks whether a segment is sorted, not by which fields — an IsSorted=true segment sorted by [pk] still passes the gate after namespaceEnabled is turned on and then gets merged by [partitionKey, pk]. That gate needs its own reachability analysis and a separate bug issue; this change only makes MergeSort fail explicitly on such input instead of panicking or corrupting.

Test gaps to close first. MergeSort currently has no varchar-key test and no multi-field-key test, and TestMergeSort's ordering assertion only checks Value(0) of each output record while the 64 MB batchSize yields a single batch. Those baselines will be added and shown green on the current code before the change lands.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area/performancePerformance issueskind/enhancementIssues or changes related to enhancementstaleindicates no udpates for 30 days

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions