Learned Indexes: Fitting a Function Instead of Storing Keys

A B-tree indexes a sorted array by storing a hierarchy of separators. Every comparison it makes narrows the search range by a factor of the node’s fanout, and the separators are keys that were computed from the data. The observation behind learned indexes is that if the data has structure, that structure can be learned and the separators need not be stored at all: a function can predict where a key belongs, and the index stores only the parameters of that function.

A dense, sorted array of n keys is a monotone function from key to rank. If the array holds 1, 2, 3, …, n, the function is a line with slope 1 and the entire index is two numbers. Real data is not a line, but it is often close to piecewise linear, and that is the regime where the idea pays.

The error bound is the contract

A learned index is a model f and a bound ε such that for every key k in the array, the true rank r(k) satisfies |f(k) − r(k)| ≤ ε. Given that guarantee, a lookup is: compute f(k), search the window [f(k) − ε, f(k) + ε] of the sorted array, and get the answer. If ε is small — say 64 — the last-mile search is a binary search over 129 elements, about 7 comparisons, or a linear scan across a few cache lines.

The bound is what makes the structure correct rather than approximate. A model with no bound is a heuristic that sometimes misses; with a bound, the structure is a genuine index with a worst-case guarantee, and the model’s accuracy determines the size of the final scan rather than whether the lookup is correct.

The consequence is that the index can be built to a target ε rather than to a target size. Tightening the bound shrinks the search window and grows the number of model segments, so a single trade-off parameter controls the space/time curve:

ε	Model segments for a smooth key set	Last-mile comparisons
1024	few	~11
64	some	~7
8	many	~4
1	potentially n	1 (direct lookup)

The last row is the ideal — a model exact enough to jump straight to the answer — and it is reachable only for data that is genuinely linear or that has been split into segments small enough to be linear.

Why one model rarely suffices

A single global model over an arbitrary sorted array can fail badly. Consider an array built by concatenating several separately sorted runs, or a key set with a sharp density change halfway through. A least-squares fit over the whole array has a residual that grows without bound in the dense region, so ε must be chosen for the worst region and becomes useless in the sparse one.

The standard answer is a hierarchy of models, and it is structurally the same as a B-tree. The top level holds a few models covering large ranges; each leaf model covers a range small enough that the local key-to-rank function is nearly linear and can be approximated within ε. The hierarchy replaces the B-tree’s separator array with model parameters, and the descent replaces comparisons with function evaluations.

The second level of the analogy matters for the cost analysis. In a B-tree, a lookup at each level does a search inside the node; in a learned index, a lookup at each level evaluates a model. Multiplication and division are slower than a comparison against a cached separator, so the learned index’s advantage is not per-level latency but total nodes touched: a model can cover a whole subtree’s range, so a hierarchy of a few models may replace three or four tree levels.

The piecewise-linear greedy construction

The practical construction is a greedy segmentation. Walk the sorted keys; maintain the segment starting at the current key; extend it as long as a line through the segment’s endpoints predicts every key’s rank within ε. When a key violates the bound, close the segment and start a new one at that key.

The check is a convex-hull condition. A line segment fitting points (k, r) with error at most ε exists exactly when the points, each expanded into a vertical interval [r − ε, r + ε], admit a line intersecting all of them. Maintaining the upper and lower hulls of the unexpanded points and checking whether the hulls’ slopes can be brought within ε gives an O(n) construction with O(1) work per key.

def segment(keys, eps):
    """Yield (start, end, slope, intercept) for each greedy segment."""
    n = len(keys)
    lo = 0
    while lo < n:
        # Fit a line through (keys[lo], lo) and extend while the bound holds.
        (x1, y1) = (keys[lo], lo)
        slope = 1.0 / max(keys[lo + 1] - keys[lo], 1) if lo + 1 < n else 1.0
        intercept = y1 - slope * x1
        hi = lo + 1
        while hi < n:
            pred = slope * keys[hi] + intercept
            if abs(pred - hi) > eps:
                break
            hi += 1
        yield (lo, hi, slope, intercept)
        lo = hi

The version above fixes the slope from the first two keys of the segment and extends until a key falls outside the bound, which is the simplest rule that produces a valid segmentation. It is not the best one: a line anchored at the segment’s start loses accuracy as the segment grows, so it closes sooner than necessary. The linear-time construction maintains the upper and lower hulls of the points and closes the segment at the exact key where no admissible line remains, which yields fewer segments for the same error bound. Both versions satisfy the same contract: every key in a segment is within ε of the line’s prediction for it, and the segments together cover the array with no gaps.

Fitting the last mile into a cache line

The reason ε in the tens rather than the thousands is useful is cache behaviour. The final search touches 2ε + 1 array entries, which for 8-byte keys is 16ε + 8 bytes. At ε = 4 that is 72 bytes, just over one 64-byte cache line; at ε = 2 it is 40 bytes and fits, though nothing forces the window to sit inside a single line — only an ε of 1 guarantees that wherever the window starts. A model accurate enough to give ε ≤ 4 turns the last mile into one or two cache-line reads and a short scan.

The scan can also be branch-free: compare all 2ε + 1 candidates with the query at once, then count the results that are not less than the query rather than the ones that are. Counting from the wrong end turns the scan into an upper bound instead of a lower bound, and the two differ by the size of the equal run whenever the key is duplicated. One comparison per candidate, summed as booleans, is a handful of instructions for small ε with no branch misprediction, and it composes with the model evaluation, which itself has no data-dependent branch.

Compare this with the structures in this cluster that optimise the same quantity. A level-order search tree reduces the number of cache lines touched by a descent to a handful by laying the top levels contiguously; a learned index reduces the number of levels by making one model cover many. The two attacks are complementary, and a learned index built over an Eytzinger-ordered array in the leaf regions is a coherent design.

Failure modes

Non-monotone keys. The construction assumes the key-to-rank function is monotone, because the rank is derived from the sort order. Duplicate keys break strict monotonicity, and a model fitted to ranks must decide whether duplicates share a rank or occupy a span. Both choices work; treating a run of duplicates as a span with a constant rank makes the function monotone non-decreasing, and the last-mile scan must then return the first index of the run — exactly the lower-bound question, with the model providing the starting point instead of the midpoint.

Distribution shift. A model is fitted to the data present at build time. If the array is later updated with keys from a different region of the key space, the model’s predictions for those keys can be arbitrarily wrong, and the error bound no longer holds. The bound is a property of the fitted data, not of the model class. This is why learned indexes are paired with a buffer for recent insertions, or are rebuilt, and why the “update” variants maintain the bound incrementally by splitting segments when accumulated error exceeds ε.

Models that encode the data. A model with enough parameters can memorise the key set rather than fit its shape, which produces ε = 0 on the training data and gives no compression. The parameter count must be compared against the separator array it replaces: if the model stores one parameter per key, it is a less convenient B-tree. The compression claim is always a claim about the ratio of model size to key count, and that ratio is what varies from one data set to another by orders of magnitude.

Floating-point rounding. A model evaluated in floating point carries its own error, and the bound must account for it. A prediction computed as slope * key + intercept in double precision has a relative error around 2⁻⁵³ per operation, which for a rank near a billion is an absolute error well below one position — small next to any useful ε, but not zero, and it accumulates across a chain of multiply-adds or when the model is evaluated in single precision. Implementations reserve a margin in the bound for it, or use fixed-point arithmetic with an explicit error term. The fixed-point route also makes the model’s evaluation bit-for-bit reproducible across platforms, which matters for an index format on disk.

Worst-case construction cost. The greedy segmentation is linear, but a model that must guarantee a bound during updates may need to split segments repeatedly, and each split requires refitting. The asymptotic guarantees of the dynamic variants come from allowing the segment count to grow, which is the same space-for-time trade a B-tree makes when it splits a node.

What the approach actually replaces

The honest framing is that a learned index replaces the separator array of a B-tree with model parameters, and replaces in-node comparison with function evaluation. It does not remove the sorted array, the need for a last-mile search, or the need for the structure to be updated.

The size claim depends on the data. For a key set generated by a monotone function with small noise, a handful of segments can cover millions of keys, and the index is three or four orders of magnitude smaller than a separator array holding every m-th key. For keys drawn from a hash — effectively random order statistics — no model predicts rank better than the identity, and the segmentation degenerates to one segment per few keys, which is larger than the separators it replaces. The technique’s applicability is a property of the data’s smoothness, and that property is measurable before committing to the design.

That is why the same idea appears in different guises across the structures here. Where the key distribution is known — a uniform distribution, as in interpolation search — a hand-derived estimate of position does the same job the fitted model does, and with the same requirement that the distribution hold. A learned index is that idea generalised to distributions the builder has not characterised, with the error bound standing in for the assumption that the model is right.