18. Reservoir Sampling for Streaming Hard Negatives

Problem

Fin's embedding model is trained with hard-negative mining — documents the reranker scored highly but that weren't actually used. In production, that candidate pool doesn't arrive as a tidy fixed-size array; it arrives as a stream of unknown, effectively unbounded length, pulled from logs as they're written. You don't know in advance how many candidates will pass by, and you can't rewind to look at one again once it's gone.

The naive approach — collect every candidate into a list, then sort or slice out the top-n — fails at this scale: it requires O(n) memory to hold the entire stream, which is exactly what you don't have when n is unbounded or the stream never truly "ends" in any useful sense. Instead, you need to maintain a fixed-size sample of k items using only O(k) memory as items flow past one at a time, updating the sample incrementally so that at every point in the stream, the k items currently held are a uniformly random sample of everything seen so far — every item that has streamed past has an equal k/n chance of being in the final reservoir, including items seen very early or very late.

Source: src/18_reservoir_sampling_hard_negatives.py

def reservoir_sample(stream: list[str], k: int, seed: int) -> list[str]:
    ...

# stream is given as a list for testability, but must be processed as a
# single forward pass — no peeking at len(stream) to shortcut the algorithm.
# seed drives a random.Random(seed) instance so results are reproducible.

Examples (exact contents depend on the seed's random stream, so these
describe the PROPERTY being demonstrated):
>>> reservoir_sample(["doc_a", "doc_b", "doc_c"], k=5, seed=0)
# k >= len(stream): every index falls in the fill phase, so the reservoir
# equals the stream exactly, in original order.
['doc_a', 'doc_b', 'doc_c']

>>> reservoir_sample(["doc_0", ..., "doc_9"], k=3, seed=42)
# len(stream) > k: reservoir holds exactly 3 items, each with an equal
# 3/10 chance of surviving — the exact winners are this seed's deterministic
# output.
['doc_4', 'doc_1', 'doc_9']

Step-by-Step Approach

1. Create an empty reservoir list and a seeded random.Random(seed) instance — all randomness must go through this instance so results are exactly reproducible.
2. Walk the stream once, tracking each item's index i via enumerate. Never consult len(stream) — the algorithm must behave as if the total length is unknown.
3. Fill phase: while i < k, append the item to the reservoir unconditionally. After this phase the reservoir holds exactly the first k items.
4. Replace phase: for every subsequent item at index i >= k, draw a random integer j uniformly from [0, i] inclusive.
5. If j < k, overwrite reservoir[j] with the current item; otherwise discard the item and move on. This is what keeps the reservoir size fixed at k without ever needing to know how much more stream is coming.
6. After the stream is exhausted, the reservoir holds the final uniformly-random sample of size k (or all n items if the stream had fewer than k elements total).

The key insight is inductive: after processing the first i+1 items, every item seen so far is in the reservoir with equal probability k/(i+1) by induction on i. When item i+1 arrives, it gets probability k/(i+1) of being chosen (via j < k), and any item already in the reservoir survives only if it isn't the one overwritten — multiplying its prior probability by 1 - 1/(i+1) · k/k = i/(i+1) preserves the k/(i+1) invariant for everyone, so even an item seen on the very first pass ends up with exactly the same k/n chance as one seen on the last.

Reference solution

import random


def reservoir_sample(stream: list[str], k: int, seed: int) -> list[str]:
    # Algorithm R: O(n) time (one pass), O(k) space (the reservoir itself).
    rng = random.Random(seed)
    reservoir: list[str] = []
    for i, item in enumerate(stream):
        if i < k:
            reservoir.append(item)
        else:
            # draw uniformly from [0, i] inclusive
            j = rng.randint(0, i)
            if j < k:
                reservoir[j] = item
    return reservoir

Key Functions & Tricks

• random.Random(seed) — instance-scoped RNG, not the global random module, so runs are isolated and reproducible
• enumerate(stream) — gives the running index i without ever consulting len(stream)
• rng.randint(0, i) — inclusive on both ends; the range grows by one each step, shrinking every item's survival odds correctly
• Fill phase (i < k) — first k items are appended unconditionally, seeding the reservoir
• Replace phase (i >= k) — every later item gets a k/(i+1) chance of bumping a random existing slot
• if j < k: reservoir[j] = item — the range check both decides acceptance and guarantees the index never runs out of bounds

How to Recognize This Pattern