pingcap / pingcap/tidb

util: add a generic, allocation-free heap to replace container/heap usage

Open
#68,161 3 comments 2 reactions 1 assignee Claimed by @hawkingrei View on GitHub
type/enhancement
Dominant language
Go
Stars
40.5k
Forks
6.2k
PR merge metrics
PR metrics pending

Description

## Enhancement

The Go standard library's `container/heap` is widely used in this repository (21+ files across planner, executor, statistics, lightning, util, etc.). Its `heap.Interface` API requires `Push(x any)` and `Pop() any`, so every push boxes the element into an `interface{}` value and every pop returns an interface that the caller type-asserts back. On hot paths this boxing dominates both wall-clock time and GC pressure.

We have two existing generic wrappers around `container/heap` in the repo:

- `pkg/util/generic.BoundedMinHeap[T]`
- `pkg/lightning/backend/external/iter.go: mergeHeap[T heapElem]`

Both improve type safety, but because they still delegate to `container/heap.Push`/`Pop`, the boxing cost is unchanged.

## Motivation: prototype benchmark

I prototyped a typed-slice heap (sift up/down on a `[]T` directly) for one k-way merge use case and compared it against the same logic using `container/heap`. Two fixtures, three partition counts, 3 benchstat runs each:

**Overlapping fixture** (heap small, pop/push count ≈ partitions × buckets-per-partition):

| Partitions | ns/op (typed prototype) | ns/op (`container/heap`) | Δ time | Δ B/op | Δ allocs/op |
|---|---|---|---|---|---|
| 100 | 11.8 ms | 15.5 ms | **+31%** | **+11×** | **+70×** |
| 1000 | 166 ms | 201 ms | **+21%** | **+12×** | **+704×** |
| 5000 | 1,133 ms | 1,319 ms | **+16%** | **+12×** | **+3,477×** |

**Skewed fixture** (each partition disjoint; pop/push count ≈ totalBuckets):

| Partitions | ns/op (typed prototype) | ns/op (`container/heap`) | Δ time | Δ B/op | Δ allocs/op |
|---|---|---|---|---|---|
| 100 | 17.4 ms | 20.3 ms | **+17%** | **+2.6×** | **+2×** |
| 1000 | 214 ms | 245 ms | **+14%** | **+2.6×** | **+2×** |
| 5000 | 1,253 ms | 1,395 ms | **+11%** | **+2.6×** | **+2×** |

The `+2×` factor for the disjoint fixture's allocs/op corresponds to exactly one extra allocation per `(push, pop)` pair vs. the typed-slice version (one box on push + one box on pop = +2 per pair). The overlapping fixture starts from a much smaller alloc baseline — the typed version barely allocates — so the relative blow-up is larger even though the absolute extra allocs per operation are similar.

## Proposal

Add a generic, allocation-free heap helper under `pkg/util/heap` (or similar):

```go
package heap

func Init[T any](h []T, less func(a, b T) bool)
func Push[T any](h *[]T, x T, less func(a, b T) bool)
func Pop[T any](h *[]T, less func(a, b T) bool) T
func Fix[T any](h []T, i int, less func(a, b T) bool)
```

Implementation is the same sift-up / sift-down algorithm as `container/heap`, but operating on a typed slice with a comparator closure — no `Push(x any)`, no `Pop() any`, no boxing. This keeps the algorithm in shared library code (no per-call-site rewrites of sift logic) while avoiding the interface-conversion cost.

## Migration scope

Audit and migrate `container/heap` users, prioritized by hot-path impact:

- **High-frequency**: planner cardinality estimation, statistics global merge, executor top-N, lightning external sort merge readers.
- **Bounded heaps**: `pkg/util/generic.BoundedMinHeap[T]` — re-implement on top of the new helper to remove boxing.
- **Ad-hoc**: any direct `container/heap` users where the heap operates in a tight loop.

Long-tail / rarely-hit users can stay on `container/heap` indefinitely; this is not a blanket replacement.

## Acceptance criteria

- New helper compiles, has unit tests covering Init / Push / Pop / Fix on a few `T` types.
- One pilot migration of a representative hot-path user, with benchstat numbers in the PR.
- The helper is documented as the preferred form for new heap code; existing `container/heap` usage is allowed to remain.

## Trade-offs

- A typed slice heap means each call site embeds the algorithm via the helper functions; binary size grows slightly with each distinct `T`.
- The comparator is passed at every call. We could alternatively define a small struct type holding `(slice, less)` but the function-argument form is simplest and keeps the API symmetrical with `container/heap`'s Init/Push/Pop shape.

## Out of scope

- Concurrent / lock-free heaps.
- Persistent / on-disk heaps.
- Replacing every `container/heap` use site (mechanical sweep is intentionally not part of this issue).

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.