db/etl: give both merges the prefix, the sift and the in-place advance

db/etl/collector.go

@@ -307,17 +307,19 @@ func mergeSortFiles(logPrefix string, providers []dataProvider, loadFunc simpleL
 		}
 	}
 
-	h := &Heap{}
-	heapInit(h)
+	h := &Heap{elems: make([]*HeapElem, 0, len(providers))}
 	for i, provider := range providers {
 		if key, value, err := provider.Next(); err == nil {
-			heapPush(h, &HeapElem{key, value, i})
+			e := &HeapElem{Value: value, TimeIdx: i}
+			e.setKey(key)
+			h.elems = append(h.elems, e)
 		} else /* we must have at least one entry per file */ {
 			eee := fmt.Errorf("%s: error reading first readers: n=%d current=%d provider=%s err=%w",
 				logPrefix, len(providers), i, provider, err)
 			panic(eee)
 		}
 	}
+	heapInit(h)
 
 	var prevK, prevV []byte
 
@@ -331,7 +333,9 @@ func mergeSortFiles(logPrefix string, providers []dataProvider, loadFunc simpleL
 			}
 		}
 
-		element := heapPop(h)
+		// The root stays in the heap while loadFunc runs, then takes its
+		// provider's next key in place: one sift instead of a pop and a push.
+		element := h.elems[0]
 		provider := providers[element.TimeIdx]
 
 		// SortableOldestAppearedBuffer must guarantee that only 1 oldest value of key will appear
@@ -363,9 +367,15 @@ func mergeSortFiles(logPrefix string, providers []dataProvider, loadFunc simpleL
 			}
 		}
 
-		if element.Key, element.Value, err = provider.Next(); err == nil {
-			heapPush(h, element)
-		} else if !errors.Is(err, io.EOF) {
+		key, value, err := provider.Next()
+		switch {
+		case err == nil:
+			element.setKey(key)
+			element.Value = value
+			heapFixRoot(h)
+		case errors.Is(err, io.EOF):
+			heapPopRoot(h)
+		default:
 			return fmt.Errorf("%s: error while reading next element from disk: %w", logPrefix, err)
 		}
 	}

db/etl/etl_test.go

@@ -2265,3 +2265,52 @@ func TestCollectRejectsOversizedKey(t *testing.T) {
 	require.Error(t, err)
 	require.Contains(t, err.Error(), "exceeds")
 }
+
+// BenchmarkProviderHeap drives mergeSortFiles' heap on its own: k sorted runs
+// of keys, advanced the way the load loop does, with no file or collector
+// around it.
+func BenchmarkProviderHeap(b *testing.B) {
+	const keyLen = 32
+	for _, k := range []int{4, 8, 16} {
+		for _, count := range []int{200_000} {
+			b.Run(fmt.Sprintf("k%d_n%d", k, count), func(b *testing.B) {
+				runs := make([][][]byte, k)
+				for i := range runs {
+					runs[i] = make([][]byte, 0, count/k)
+					for j := range count / k {
+						key := make([]byte, keyLen)
+						x := uint64(i*count+j) * 6364136223846793005 //nolint:gosec
+						binary.BigEndian.PutUint64(key, x)
+						binary.BigEndian.PutUint64(key[8:], x^0xdeadbeef)
+						runs[i] = append(runs[i], key)
+					}
+					slices.SortFunc(runs[i], bytes.Compare)
+				}
+				val := make([]byte, 128)
+				b.ResetTimer()
+				for b.Loop() {
+					at := make([]int, k)
+					h := &Heap{elems: make([]*HeapElem, 0, k)}
+					for i := range runs {
+						e := &HeapElem{Value: val, TimeIdx: i}
+						e.setKey(runs[i][0])
+						h.elems = append(h.elems, e)
+						at[i] = 1
+					}
+					heapInit(h)
+					for h.Len() > 0 {
+						e := h.elems[0]
+						i := e.TimeIdx
+						if at[i] < len(runs[i]) {
+							e.setKey(runs[i][at[i]])
+							at[i]++
+							heapFixRoot(h)
+							continue
+						}
+						heapPopRoot(h)
+					}
+				}
+			})
+		}
+	}
+}

db/etl/heap.go

@@ -18,116 +18,109 @@ package etl
 
 import (
 	"bytes"
+	"encoding/binary"
 	"slices"
 )
 
+// HeapElem is the entry a provider currently sits on. pfx caches the first 8
+// key bytes, so a comparison reaches Key only when two providers agree there.
 type HeapElem struct {
+	pfx     uint64
 	Key     []byte
 	Value   []byte
 	TimeIdx int
 }
 
+func (e *HeapElem) setKey(k []byte) { e.Key, e.pfx = k, keyPrefix(k) }
+
+// Heap orders providers by the key each one sits on. Providers are numbered in
+// the order their files were written, so the lower TimeIdx wins a tie and
+// equal keys come out in the order they went in.
 type Heap struct {
 	elems []*HeapElem
 }
 
-func (h *Heap) Len() int {
-	return len(h.elems)
-}
+func (h *Heap) Len() int { return len(h.elems) }
 
-func (h *Heap) Less(i, j int) bool {
-	if c := bytes.Compare(h.elems[i].Key, h.elems[j].Key); c != 0 {
+func lessElem(a, b *HeapElem) bool {
+	if a.pfx != b.pfx {
+		return a.pfx < b.pfx
+	}
+	if c := bytes.Compare(a.Key, b.Key); c != 0 {
 		return c < 0
 	}
-	return h.elems[i].TimeIdx < h.elems[j].TimeIdx
-}
-
-func (h *Heap) Swap(i, j int) {
-	h.elems[i], h.elems[j] = h.elems[j], h.elems[i]
-}
-
-func (h *Heap) Push(x *HeapElem) {
-	h.elems = append(h.elems, x)
-}
-
-func (h *Heap) Pop() *HeapElem {
-	old := h.elems
-	n := len(old) - 1
-	x := old[n]
-	//old[n].Key, old[n].Value, old[n].TimeIdx = nil, nil, 0
-	old[n] = nil
-	h.elems = old[0:n]
-	return x
+	return a.TimeIdx < b.TimeIdx
 }
 
-// ------ Copy-Paste of `container/heap/heap.go` without interface conversion
-
-// Init establishes the heap invariants required by the other routines in this package.
-// Init is idempotent with respect to the heap invariants
-// and may be called whenever the heap invariants may have been invalidated.
-// The complexity is O(n) where n = h.Len().
+// heapInit orders elems appended in provider order.
 func heapInit(h *Heap) {
-	// heapify
-	n := h.Len()
-	for i := n/2 - 1; i >= 0; i-- {
-		down(h, i, n)
+	for i := len(h.elems)/2 - 1; i >= 0; i-- {
+		h.siftRoot(i)
 	}
 }
 
-// Push pushes the element x onto the heap.
-// The complexity is O(log n) where n = h.Len().
-func heapPush(h *Heap, x *HeapElem) {
-	h.Push(x)
-	up(h, h.Len()-1)
-}
-
-// Pop removes and returns the minimum element (according to Less) from the heap.
-// The complexity is O(log n) where n = h.Len().
-// Pop is equivalent to Remove(h, 0).
-func heapPop(h *Heap) *HeapElem {
-	n := h.Len() - 1
-	h.Swap(0, n)
-	down(h, 0, n)
-	return h.Pop()
-}
+// heapFixRoot restores the order after the root provider moved to its next key.
+func heapFixRoot(h *Heap) { h.siftRoot(0) }
 
-func up(h *Heap, j int) {
-	for {
-		i := (j - 1) / 2 // parent
-		if i == j || !h.Less(j, i) {
-			break
-		}
-		h.Swap(i, j)
-		j = i
+// heapPopRoot drops the root, whose provider has no more keys.
+func heapPopRoot(h *Heap) {
+	last := len(h.elems) - 1
+	h.elems[0] = h.elems[last]
+	h.elems[last] = nil
+	h.elems = h.elems[:last]
+	if last > 0 {
+		h.siftRoot(0)
 	}
 }
 
-func down(h *Heap, i0, n int) bool {
-	i := i0
+// siftRoot restores the heap under i, whose element changed: sink the hole to
+// a leaf taking the smaller child, then climb back until the old element fits.
+// One compare a level rather than the two a top-down sift needs, and the
+// element that just moved on usually holds a larger key now, so the hole
+// nearly always reaches a leaf. The climb back is also the heapify step.
+func (h *Heap) siftRoot(i int) {
+	x, top, n := h.elems[i], i, len(h.elems)
 	for {
-		j1 := 2*i + 1
-		if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
+		l := 2*i + 1
+		if l >= n {
 			break
 		}
-		j := j1 // left child
-		if j2 := j1 + 1; j2 < n && h.Less(j2, j1) {
-			j = j2 // = 2*i + 2  // right child
+		if r := l + 1; r < n && lessElem(h.elems[r], h.elems[l]) {
+			l = r
 		}
-		if !h.Less(j, i) {
+		h.elems[i] = h.elems[l]
+		i = l
+	}
+	for i > top {
+		p := (i - 1) / 2
+		if lessElem(h.elems[p], x) {
 			break
 		}
-		h.Swap(i, j)
-		i = j
+		h.elems[i] = h.elems[p]
+		i = p
 	}
-	return i > i0
+	h.elems[i] = x
 }
 
 // ------ the merge over a sortableBuffer's sorted chunks
 
+// keyPrefix is the first 8 bytes of k, zero-padded. Big-endian, so comparing
+// two prefixes as integers orders them the way bytes.Compare would.
+func keyPrefix(k []byte) uint64 {
+	if len(k) >= 8 {
+		return binary.BigEndian.Uint64(k)
+	}
+	var pad [8]byte
+	copy(pad[:], k)
+	return binary.BigEndian.Uint64(pad[:])
+}
+
 // merger walks already-sorted chunks in key order, a cursor per chunk under a
-// heap of chunk ids.
+// heap of chunk ids. The prefixes live apart from the cursors because the heap
+// reads them on nearly every comparison and they stay in L1.
 type merger struct {
 	heap []int32  // chunk ids, ordered by their cursor's key
+	pfx  []uint64 // each cursor's key prefix, by chunk id
 	cur  []cursor // by chunk id
 
 	// Chunks already in order end to end, which ascending keys produce, are
@@ -147,6 +140,7 @@ type cursor struct {
 func (m *merger) rewind(chunks []dataChunk) {
 	clear(m.cur) // a shorter run would leave the old cursors pinning their chunks
 	m.cur = slices.Grow(m.cur[:0], len(chunks))[:len(chunks)]
+	m.pfx = slices.Grow(m.pfx[:0], len(chunks))[:len(chunks)]
 	for i := range chunks {
 		m.cur[i] = cursor{ents: chunks[i].entries(), buf: chunks[i].buf}
 	}
@@ -204,13 +198,14 @@ func (m *merger) next() ([]byte, entryLoc, bool) {
 
 func (m *merger) release() {
 	clear(m.cur)
-	m.cur, m.heap = m.cur[:0], m.heap[:0]
+	m.cur, m.pfx, m.heap = m.cur[:0], m.pfx[:0], m.heap[:0]
 	m.chunk, m.concat = 0, false
 }
 
 func (m *merger) load(id int32) {
 	c := &m.cur[id]
 	c.key = keyOf(c.buf, c.ents[c.at])
+	m.pfx[id] = keyPrefix(c.key)
 }
 
 // chunksInOrder reports whether every chunk's last key comes before the next
@@ -236,26 +231,40 @@ func (m *merger) chunksInOrder() bool {
 // less orders two cursors by the key they sit on. Chunks fill in insertion
 // order, so the lower id wins a tie and equal keys keep the order they went in.
 func (m *merger) less(x, y int32) bool {
+	if px, py := m.pfx[x], m.pfx[y]; px != py {
+		return px < py
+	}
 	if r := bytes.Compare(m.cur[x].key, m.cur[y].key); r != 0 {
 		return r < 0
 	}
 	return x < y
 }
 
-// siftRoot restores the heap under i, whose element changed.
+// siftRoot restores the heap under i, whose cursor just moved: sink the hole
+// to a leaf taking the smaller child, then climb back until the old value
+// fits. One compare a level instead of the usual two, and the cursor that just
+// won usually holds a larger key, so the hole nearly always reaches a leaf.
+// The climb back is also what makes it serve as the heapify step.
 func (m *merger) siftRoot(i int) {
+	x, top := m.heap[i], i
 	for {
-		s, l, r := i, 2*i+1, 2*i+2
-		if l < len(m.heap) && m.less(m.heap[l], m.heap[s]) {
-			s = l
+		l := 2*i + 1
+		if l >= len(m.heap) {
+			break
 		}
-		if r < len(m.heap) && m.less(m.heap[r], m.heap[s]) {
-			s = r
+		if r := l + 1; r < len(m.heap) && m.less(m.heap[r], m.heap[l]) {
+			l = r
 		}
-		if s == i {
-			return
+		m.heap[i] = m.heap[l]
+		i = l
+	}
+	for i > top {
+		p := (i - 1) / 2
+		if m.less(m.heap[p], x) {
+			break
 		}
-		m.heap[i], m.heap[s] = m.heap[s], m.heap[i]
-		i = s
+		m.heap[i] = m.heap[p]
+		i = p
 	}
+	m.heap[i] = x
 }