Reservoir sampling

Sometimes you need a random sample of a stream, such as 1,000 log lines from today's traffic, or 100 requests to replay in a load test. The stream may be huge, and you do not know its length until it ends. You cannot pick a random index first, because you do not know the range. Reservoir sampling solves this: one pass, and memory for only k items.

The big idea

Keep a list of k items, the reservoir. Fill it with the first k items. After that, handle each new item i like this:

• With probability k / i, the new item enters the reservoir and replaces a random existing item.
• Otherwise, skip it.

That is the whole algorithm. The probability shrinks as the stream grows, which keeps the sample fair across the whole stream, not just the start.

Walking through a small run

Take a stream of 10 items and a reservoir of 3. Here is the first step after the reservoir is full.

after items 1, 2, 3   reservoir = [1, 2, 3]
item 4 arrives:       chance to enter = k / i = 3 / 4 = 0.75
   pick j = random in 0..3 (4 values)
   if j < 3, replace reservoir[j] with 4
   otherwise, skip item 4

Item 4 enters with probability 3/4, and it replaces one of the three slots. Each later item follows the same rule, using its own position i.

Code

import random

def reservoir_sample(stream, k):
    reservoir = []
    for i, item in enumerate(stream, start=1):
        if i <= k:
            reservoir.append(item)            # fill the reservoir first
        else:
            j = random.randrange(i)           # random number in 0..i-1
            if j < k:                         # happens with probability k / i
                reservoir[j] = item
    return reservoir

Check that every item has the same chance. Run it 20,000 times on a stream of 10 items, with a reservoir of 3:

from collections import Counter

random.seed(1)
counts = Counter()
for _ in range(20_000):
    counts.update(reservoir_sample(range(10), 3))
print({k: round(v / 20_000, 3) for k, v in sorted(counts.items())})

Output:

{0: 0.3, 1: 0.299, 2: 0.303, 3: 0.295, 4: 0.293, 5: 0.302, 6: 0.303, 7: 0.305, 8: 0.301, 9: 0.299}

Every item lands in the sample about 30% of the time, which is k / n = 3 / 10. Item 0 has no advantage over item 9.

Why every item is equally likely

Track one item, number i. It enters the reservoir with probability k / i. After that, it survives the next item, i + 1, unless that item picks its slot.

P(item i survives item i+1) = 1 − (k / (i+1)) × (1 / k) = i / (i+1)

The item survives the next step if the new item does not pick its slot. Multiply the survival chances from i + 1 to the end of a stream of length n, and the product collapses:

P(item i in final sample) = (k/i) × (i/(i+1)) × ((i+1)/(i+2)) × ... × ((n−1)/n)
                          = k / n

Each i cancels with the next term's numerator, so every item ends with probability k / n.

Variants and merging

• Weighted sampling gives heavier items a better chance. Efraimidis and Spirakis' A-Res method keeps the k items with the largest random()^(1/weight).
• Skipping ahead. For very long streams, Algorithm L draws how many items to skip, so most items are never examined.
• Shards. Reservoirs from different shards can be combined if you track how many items each shard saw. The count is what makes the combined sample fair.

Common mistakes

• Taking the first k items. That sample favors the start of the stream.
• Sampling with a fixed probability like 1%. You get a random number of items, not exactly k.
• Forgetting the count when combining shards. Without each shard's count, the combined sample can be biased.

Check yourself

1. Why does the probability k / i shrink as the stream grows? Each new item competes with more earlier items. The reservoir stays a fixed size, so a later item gets a smaller chance to enter.
2. A stream of 1,000,000 items is sampled with k = 100. What chance does any one item have? 100 / 1,000,000 = 0.01%.

Key takeaways

• Reservoir sampling draws a uniform sample of k items from a stream of unknown length in one pass.
• Fill the reservoir first. Then accept item i with probability k / i.
• Every item ends with probability k / n. The telescoping product proves it.
• Keep each shard's count if you merge samples from several shards.

Further learning

Next, read MinHash and LSH for estimating how similar two sets are, which builds on the same idea of using a random choice to get a fair estimate.