Imagine you are processing a stream of events, of unknown length. It could end in 3 seconds, it could run for 3 months; you simply don’t know. As a result, storing the whole stream in memory or even on disk is not acceptable, but you still need to extract relevant information.

Depending on what information you need, choosing a random sample of the stream will give you almost as good information as storing all the data. For example, consider a performance profiler, used to find which parts of your running code are slowest. Many profilers records a program’s callstack every few microseconds, resulting a stream of unlimited size: you don’t know how long the program will run. For this use case, a random sample of callstacks, say 2000 of them, can usually give you sufficient information to do performance optimization.

Why does this work?

• Slow code will result in the same callstack being repeated.
• A random sample of callstacks is more likely to contain callstacks that repeat a lot.
• Thus, a random sample is more likely to include slow code, the code you specifically want to identify with your profiler.

When you need to extract a random sample from a stream of unknown length, a common solution is the family of algorithms known as reservoir sampling. In this article you will learn:

• How basic reservoir sampling works.
• Some problems with reservoir sampling, motivated by a profiler that wants to generate a timeline.
• A (new?) variant of reservoir sampling that allows you to ensure samples are spread evenly across time.

Let’s start with the simplest case: you want to choose a single sample from an event stream of unknown length. The basic algorithm:

1. Record the first event in the stream.
2. When you reach the i-th event, swap it for the currently chosen event with a chance of 1/i.

Why does this work? Imagine you have a stream of N events; each event should have a 1/N chance of being chosen.

• The last, Nth event is chosen 1/Nth the time, which is what we want.
• The N-1 item is swapped in 1/(N-1) of the time, and then has a (N-1)/N chance of being kept in the final swap. Which means it has a 1/(N-1) * (N-1)/N = 1/N chance of being chosen.
• The N-2 item is swapped in 1/(N-2) of the time, and has a (N-2)/(N-1) * (N-1)/N chance of being swapped in.
• Etc.

Here’s what this looks like in Python:

from random import randrange

def reservoir1(events):
    result = None
    samples = 0
    for event in events:
        samples += 1
        choice = randrange(0, samples)
        if choice == 0:
            result = event
    return result

Note that there are more efficient algorithms, I am demonstrating this one because it’s so simple.

Now that you’ve seen how to pick one item out of a stream, picking K items is an extension of the same basic idea:

from random import randrange

def reservoir(events, num_items):
    events = iter(events)
    # Until we have enough items, we pick them all:
    result = [next(events) for _ in range(num_items)]
    samples = num_items
    for event in events:
        samples += 1
        choice = randrange(0, samples)
        if choice < num_items:
            # Replace that item:
            result[choice] = event
    return sorted(result)

Let’s return to the example of a profiler. If you pick K samples at random, and K is sufficiently large, you will be able to identify the slowest parts of the code, since slow code is more likely to show up in the samples. You can shove the resulting samples into a flamegraph visualization, and you will get a reasonable visualization of the callstacks of the slow code. And when profiling a high-volume server, with uniform behavior over time and no clear beginning or end, this works fine.

But now let’s say you’re profiling a batch job, for data science or scientific computing. Here there’s a start, middle, and end, each doing different things: loading data, processing data, storing results. There may be different thread pools or concurrency approaches used in different parts of the program. In this situation, flamegraphs are less useful:

• You care about ordering in time, but flamegraphs have no ordering.
• You care about concurrency, but flamegraphs can’t show concurrency very well.

As a result, instead of or in addition to a flamegraph, you may use a timeline visualization, showing what happened over time.

But now you have a new problem: with reservoir sampling, any given sample doesn’t represent a well-defined amount of time. If you have 100 consecutive events (0 through 99) from which you are picking 5 random samples, you might conceivably get events 1, 3, 17, 72, and 95 sampled; because the gaps vary so much, visualizing each of these five event as covering 20% of the timeline would be misleading. And giving them differing amounts of time in the timeline can also be misleading: just because there’s a large gap between 17 and 72 doesn’t mean that those sampled callstacks lasted longer than callstack 3.

We can demonstrate the uneven gaps visually. Given a series of 1000 events, and reservoir sampling that samples 10 events, here’s the histogram of the difference in event index between two subsequent sampled events. For example, a difference of 5 means that 5 events elapsed between one event being sampled and the next being sampled.

Notice that the gap between two samples is often quite small, and sometimes can be quite large. If you constructed a timeline based on these samples, and represented each sample as being the same amount of time, you’d be distorting the amounts of time involved.

To solve this problem, I’m going to demonstrate what I call “timesliced reservoir sampling”. Most likely I’m just reinventing something that already exists, email me if you know of prior art.

In normal reservoir sampling, you choose K samples from the stream. In timesliced reservoir sampling, you divide the stream into K equal-sized timeslices, and then pick one random sample from each timeslice using reservoir sampling. That means the samples will be much more evenly spread out. Of course, you don’t know how long the stream is, so dividing it exactly into K timeslices is impossible. But you do know how many events you’ve gotten so far, which allows for an approximation that can be updated as more events arrive.

To make implementation easier, this algorithm relaxes two requirement:

• Instead of returning exactly K samples, it will return between K and 2K samples. For a profiler, this doesn’t meaningfully impact the visualization for users.
• It will also not be quite fair in the last sample; once K is large enough this doesn’t meaningfully distort the result.

Here’s an implementation of the algorithm:

from random import choice

def take(iterator, n):
    """Yield the first N items of an iterator."""
    for i, value in enumerate(iterator):
        yield value
        if i + 1 == n:
            return

def timesliced_reservoir(events, num_timeslices):
    timeslice_size = 1
    result = []
    events = iter(events)
    while True:
        next_batch = take(events, timeslice_size)
        next_event = reservoir1(next_batch)
        if next_event is None:
            break
        result.append(next_event)

        if len(result) == 2 * num_timeslices:
            # Our current estimate of timeslice size has
            # grown:
            timeslice_size *= 2
            # Cut result's size in half, by dividing the
            # current result into pairs, and randomly
            # picking one event from each pair:
            result = [
                choice(result[2 * i: 2 * (i + 1)])
                for i in range(num_timeslices)
            ]
    return result

Once more we can visualize the gap between consecutive samples. This time the gaps are much more consistent:

Start with a timeslice size T = 1, and an empty list result.

Then, as long as there are remaining items:

1. Use reservoir sampling to pick on item out of the next T items.
2. Add that item to the result.
3. If the result is now 2K long: double T, and reduce down to K samples by choosing one item from each consecutive pair.
4. Repeat.

Why this works is somewhat easier to understand if you track the internal state of T and result. In the following example, we have a stream of 50 events (0, 1, 2, … 49), and the algorithm was told to pick K = 4 samples. That means it will return between 4 and 8 samples. The internal state of T and result is recorded every time result changes, i.e. every time we pick a new item with normal single-item reservoir sampling from the next batch of T items.

Notice that every time result’s length reaches 8, its size is cut in half by picking one event out of each consecutive pair, and the estimated timeslice size T is doubled:

T  |  result
===|========
1  |  []
1  |  [0]
1  |  [0, 1]
1  |  [0, 1, 2]
1  |  [0, 1, 2, 3]
1  |  [0, 1, 2, 3, 4]
1  |  [0, 1, 2, 3, 4, 5]
1  |  [0, 1, 2, 3, 4, 5, 6]
1  |  [0, 1, 2, 3, 4, 5, 6, 7]
2  |  [1, 2, 4, 7]
2  |  [1, 2, 4, 7, 9]
2  |  [1, 2, 4, 7, 9, 11]
2  |  [1, 2, 4, 7, 9, 11, 12]
2  |  [1, 2, 4, 7, 9, 11, 12, 14]
4  |  [2, 4, 9, 14]
4  |  [2, 4, 9, 14, 18]
4  |  [2, 4, 9, 14, 18, 22]
4  |  [2, 4, 9, 14, 18, 22, 27]
4  |  [2, 4, 9, 14, 18, 22, 27, 31]
8  |  [4, 9, 22, 31]
8  |  [4, 9, 22, 31, 32]
8  |  [4, 9, 22, 31, 32, 45]
8  |  [4, 9, 22, 31, 32, 45, 49]

In this case, the algorithm returns 7 samples, more or less evenly spread out across 7 timeslices. Since T = 8 when the algorithm finished, it had decided that each timeslice was approximately 8 events long, and indeed 8 * 7 = 56, which is approximately 50.

Some notes on real-world usage:

• Insofar as timesliced reservoir sampling builds on standard single-item reservoir sampling, you can get faster results by using a faster algorithm for that.
• For a profiler, you will typically be sampling callstacks from multiple threads. You’ll want to align the randomly-generated choices across threads so that the reported callstacks line up accurately across threads. So for example you want to sample events 1, 3, and 7 from all threads, rather than one thread being events 1, 3, 7 and another 2, 5, 6.
• Rather than collecting all callstacks and then using the reservoir sampling to decide which to keep them, you can do the reservoir sampling first and then use that to decide whether there is any need to collect callstacks.

As mentioned, it’s not clear to me if timesliced reservoir sampling is already documented somewhere else. If you’ve seen it or used a variant of this before, please let me know. I’d also love to know if you end up using this algorithm.

Update: I have found one similar algorithm in the wild, so I imagine there are others.