For most operating systems, when memory usage grows beyond what the hardware can support, the kernel starts swapping, i.e. moving data from memory to disk. In this post, we’ll study how the Linux kernel does this, as part of the more general memory reclamation process. We’ll also cover the memory allocation flow, because memory reclamation is just one part of it, one of the unhappy paths.

First, we’ll go over the different components involved, and then how they fit together in the memory allocation flow and then memory reclamation via swap and zswap.

Concepts

Pages

A page is a chunk of contiguous memory. We’ll assume an x86-64 Linux configuration and that a page is 4KB ($2^{12}$ bytes). A page can be virtual (a chunk of virtual memory) or physical (a chunk of physical memory). To avoid ambiguity, we’ll assume a page is virtual and denote a physical page as a page frame.

Pages are useful because that’s the minimum granularity the kernel supports in terms of allocation.

Virtual Memory

When a process is created, it gets its own virtual address space. In 64-bit systems it can theoretically represent $2^{64}$ bytes, but in practice the range is smaller due to how addresses are encoded.

We’ll talk about page tables later, but it suffices to know a virtual address is composed of 4-5 regions + offset. The regions identify the location of the page, and the offset is the distance from the start of that page. Each region supports 9 bits, and since a page has 4KB, we need 12 bits to address it, so a system with 5 regions can represent $5 \cdot 9 + 12 = 57$-bit virtual addresses.

When we use syscalls like mmap() or brk(), the kernel reserves a range of the virtual address space, called a Virtual Memory Area (VMA), and returns the start of that virtual address range. This address is not backed by physical memory.

The list of VMAs is kept in a data structure called Maple tree. We won’t go over details, but it can efficiently determine the VMA of a given virtual address. We’ll call this structure the VMA tree.

There’s a special section of a process’ virtual address space which is reserved for the kernel, usually high addresses. The reason for this is so that kernel objects don’t need a special address space. They can use the same space as the process being executed.

Physical Memory

The physical memory is the RAM that you can actually write to. As we said, virtual memory is not initially backed by physical memory. This only happens when we write to it, which triggers a page fault, which causes the kernel to finally try to back the page being written to by a page frame. We’ll cover this flow in more detail later. The important part to note is that mapping virtual pages to physical pages is lazy.

Another important fact is that the kernel overcommits (analogous to overbooking in airlines). When multiple callers call mmap(), it won’t fail the call if the amount of virtual memory used exceeds the host or cgroup memory. That’s because the kernel knows that most of the time not all the requested memory will be used (think of .reserve() vs .size() in std::vector).

Recall that two different virtual addresses can map to the same physical one. For example, when we fork() a process, the child process inherits the addresses from the parent, including the mapping to physical memory. However, this is COW (copy-on-write): the moment one attempts to write, the contents are copied to a new location.

As we mentioned, the physical page is also 4KB and aligned to 4KB, so we can drop the 12 least significant bits when representing it. This representation is called a page frame number or PFN. The kernel maintains a structure with metadata on a struct page (ignore the contents for now):

struct page {
    /* ... */
};

Conceptually, this struct lives in a giant array, vmemmap, and the PFN is an index in this array. So if we want to access the metadata for a given PFN, we do vmemmap + PFN.

Note that because this is a struct that lives in the kernel, it’s subject to the same memory layout: it’s divided into virtual pages that are mapped to physical pages, which have metadata in vmemmap. There’s some degree of circularity that can be mind-bending.

Page Table

The mapping from virtual to physical memory is kept in a structure called a page table. Conceptually, it’s a map from the 57-bit virtual address (see Virtual Memory) to a 64-bit number, encoding the physical address. Because we can drop 12 bits of the physical address (due to 4KB alignment), we can reuse them for metadata / flags.

The problem with storing all virtual page mappings in a single contiguous array is that this would have to be mapped to a similarly gigantic array in the physical memory. Instead, we use an n-ary tree where each node is a page. The “regions” we mentioned in Virtual Memory are each level of this tree, so for 5 regions the tree has depth 5.

Each internal node has 512 entries pointing to the 64-bit physical address of the children nodes. Each leaf node contains 512 entries of the 64-bit physical address, which are the mappings for the virtual addresses. But which virtual address? Recall that a virtual address has 5 parts that identify the virtual page + an offset. Each of these parts corresponds to the entry index in a given node.

For example, if we have the parts 41, 121, 5, 203, 57, this tells us: in the root, get the physical address at entry 41 and navigate to that child. Then do the same for entry 121 and so on. After following 203, we’ll arrive at the appropriate leaf. Now we look up the entry 57 and the corresponding physical address is the mapping for the virtual address 41, 121, 5, 203, 57! So 41, 121, 5, 203, 57 is really a path on this tree (think of JSON paths).

The regions or levels in this tree have weird acronyms: PGD (root, Page Global Directory), P4D (level 1, Page 4th-level Directory), PUD (Page Upper Directory), PMD (Page Middle Directory), PTE (leaf, Page Table Entry). There’s one such structure per virtual address space, so simplistically, one per process.

Note that this structure grows on demand. If there’s only one entry, we’d have one node per region. When the 512 entries on the first leaf fill up, we add a new leaf, etc.

Note that the page table works directly with physical addresses, not virtual. This is very important because the CPU (see TLB) must be able to traverse this table and it has no knowledge of virtual addresses.

Kernel. As we mentioned in Virtual memory, a portion of the process’ virtual address space is reserved for the kernel. Since the mapping is consistent across processes, they can actually share the same “node” in the page table. In Figure 1, the yellow page contains addresses for the kernel virtual spaces. So different processes would have their own page table, but they would “point” to this same yellow node.

Bootstrapping. Note that the nodes themselves are pages, so their addresses must be in the page table! In Figure 1, the yellow node contains the mapping for the other nodes. For example, 15-15-1 maps to 5831047296.

This creates an interesting problem: how are the first nodes added to this tree? When the kernel is bootstrapped, it can initialize a minimal table with some pages pre-mapped.

Zones

Zones are partitions of the physical address space. Typically, there are 3 zones: ZONE_DMA for low addresses (less than 16MB), ZONE_DMA32 for the 32-bit address space (less than 4GB), and ZONE_NORMAL for the rest. These exist for legacy reasons; for example, ZONE_DMA32 is needed to support devices limited to 32-bit addresses. For hardware using NUMA, there’s at least one zone per NUMA node.

The partition of addresses to zones is fixed: we cannot transfer pages between them. It is possible to copy pages between zones and the kernel might do it outside of the memory reclamation process.

Buddy Allocation

The Page Table stores the mapping from virtual to physical, but how does it decide which physical page to map to? The kernel uses an implementation of the Buddy Memory Allocation, which is also used by the Jemalloc Memory Allocator.

We don’t need to go into the details of the algorithm, except that each zone has its own buddy allocator and its “domain” is pretty much the entire physical address space of that zone. The minimum granularity of the allocator is 4KB, i.e., one physical page, which is typically what the page fault trigger requests. The kernel might ask for larger chunks for other purposes.

The buddy allocator knows how much free memory it has.

TLB

TLB stands for Translation Lookaside Buffer, and it’s a cache that exists in the CPU. When the CPU executes instructions, the addresses it sees in registers are in virtual space. So if it’s to read and write data to/from memory, it needs to resolve the physical address.

This is an expensive operation, so it caches entries in the TLB. We mentioned TLB in the context of CPU Cache. If it gets a cache miss, it needs to traverse the page table from the PGD node, which can be done with 5 lookups (one per level). Note that this process doesn’t involve the kernel.

Folio

A folio is a sequence of contiguous physical pages. The crazy thing is that a struct folio is compatible with a struct page. For example, suppose PFN 10 is the head of a sequence of physical pages. Then if we do:

void *X = (void *)(vmemmap + 10);
struct page *p = (struct page *)X;
struct folio *f = (struct folio *)X;

Here p represents a struct page, but the bytes in memory can also represent a struct folio.

MGLRU

MGLRU stands for Multi-Generational Least Recently Used. It’s a structure that can tell which mapped pages are more suitable to be swapped. It’s similar to an LRU cache, but it cannot afford to keep the LRU order up-to-date like a user-space structure such as CacheLib can. It actually stores folios (a run of pages), not individual pages.

It uses an interesting algorithm based on generations. Whenever a page is accessed, the CPU sets a bit in the corresponding entry in the PTE (recall it has 12 bits for metadata).

Aging. Conceptually, when a request comes for reclaiming pages, it might decide to create a new generation node which contains a linked list. It scans all the pages and adds the ones with the bit set in the PTE to the list, setting the bit back to 0. The implementation is more clever: the struct folio has fields for making it a node in a doubly linked list, i.e., prev and next pointers to other folios (which is also what CacheLib does), so it builds the list by changing these pointers.

Then the generation node points to the first “node”. We don’t create a new node every time a new generation is created: this structure is a circular queue or a ring buffer.

kswapd

This is the kernel swap daemon. It’s actually an OS thread that sleeps until memory pressure wakes it up (see Watermark). It then tries to reclaim memory by asking the MGLRU to give it a candidate.

A reason I was given for it being a background thread that is idle / sleeping most of the time instead of being created on demand is that it’s mostly useful during memory pressure, and having to create a thread and allocate resources at that time is risky.

Watermark

The watermarks are actually thresholds, not like “highest value we’ve seen so far” as in stream processing. There are three watermarks: high, low, min.

If the amount of free memory (see Buddy Allocation) drops below the low threshold, it wakes up the kswapd which tries to reclaim memory until the level goes above high (this is similar to the 2-threshold approach used by thermostats to avoid flapping on/off). This is the asynchronous path. On the other hand, if a thread would cause free memory to fall below min, it will be synchronously required to perform memory reclamation. We’ll cover this synchronous path as an example flow later.

cgroups

So far we’ve only considered memory at a global level. However, cgroups are taken into account as well. At a high level, MGLRU has knowledge about cgroups, so memory pressure at individual cgroups can be taken into account to decide which pages to evict. So, for example, if we had two candidate pages that haven’t been accessed recently, we could look at their cgroup stats to make a decision.

Cgroups have the so-called interface files which are listed under /sys/fs/cgroup/<my_cgroup_path>/, and many of them are memory-related, such as memory.current, memory.min, memory.low, memory.high and memory.max. The easiest is memory.current, which corresponds to how much memory a given cgroup is using.

The file memory.min indicates the threshold for which memory is protected: the kernel will not try to reclaim memory from that cgroup if the resulting memory.current dips below that value. The threshold memory.low is similar, but it’s a soft limit: the kernel will avoid reclaiming memory for that cgroup, but if it can’t find any candidate cgroups, it will still reclaim memory.

Memory reclamation can happen at the cgroup level as well. When a task/thread for a given process causes memory.current >= memory.high, this thread is synchronously required to reclaim memory. The thread might also be explicitly put to sleep (throttled) as a backpressure mechanism. If memory.current >= memory.max and the thread cannot bring it down via reclamation, the OOM killer kills either the requesting process or the entire cgroup (depending on the memory.oom.group setting).

Flows

Now that we covered the major components involved in memory reclamation, we can focus on a specific flow, starting with a page fault.

Page Fault

As we discussed in TLB, when the CPU tries to execute an instruction accessing memory, it needs to resolve the physical address. It will first look in the TLB, then traverse the page table for that address.

If it doesn’t find it and it’s a read, it will raise a page fault exception. The kernel has a fault handler that then gets executed. The first thing it does is to determine if this virtual address is part of a valid VMA by looking it up in the VMA tree (see Virtual memory). If not, it throws a SIGSEGV (segmentation fault). Sometimes this is displayed as:

code: address not mapped to object.

This can be misleading because “address not mapped to object” is essentially why a page fault happens, and that’s technically the “happy case”: the fault handler is expected to run. In a SIGSEGV, though, it won’t because this is not a valid VMA.

If it is a valid address, then it’s now that the mapping happens. If this is from user-space code, most likely the kernel will choose ZONE_NORMAL. The corresponding buddy allocator will attempt to give it a physical page.

If it succeeds, the page fault handler will add a new entry to the page table, based on the virtual address. It adds “nodes” to the table if they don’t exist already and returns. The CPU will retry the instruction and succeed this time.

Now let’s cover the scenario in which the amount of free physical memory falls below the low watermark.

Memory Reclamation

The first thing the kernel will do is to wake up the kswapd thread. The kernel will continue running and potentially executing other instructions.

As we mentioned, the kswapd will check MGLRU, which will look for its oldest generations. If it finds pages/folios there, it gives them back to kswapd. It might also create a new generation as part of this process (see Aging in MGLRU).

If this page/folio is file-backed (i.e., a cache of a file in memory), it first determines if it’s clean (i.e., it hasn’t been modified in memory). If not, it first needs to write it back to the file. If it’s clean, it can just remove the page and the necessary references. We will not cover this flow here.

Suppose the page/folio is not file-backed. If the system does not have swap enabled, it might invoke the OOM killer. Let’s consider the case in which swap is enabled.

Swap

There are 2 sinks of swap: a partition and a file. A partition is a region of disk/flash that is dedicated for swap. Writing to it does not involve the file system. A file-based swap involves writing the page to the partition dedicated to the general filesystem. In practice, for modern Linux kernels the difference between these two is not very substantial, however. This process of writing to disk is called swap out.

Once a page is successfully copied to the appropriate swap region/device, the page is “returned” back to the corresponding buddy allocator. It must also be evicted from the TLB of every CPU core that might have it.

The swap is identified by the device type plus an offset. This identifier is then used in the page table, but the page is marked as not present, so if the CPU tries to read/write to this page, a page fault occurs. It then goes through the usual flow of requesting a physical page from the buddy allocator or trying to evict an existing page. Once it gets a new physical address, it writes to the page table and copies the contents from disk to that page; this process is known as swap in.

If every time a page is swapped in, it requires swapping out another page, the system is essentially using the disk as its RAM and performance degrades. This is called thrashing.

Zswap

At a high level, zswap intercepts pages about to be swapped and instead compresses them using encoders such as zstd (see LZ77) and keeps them in memory instead. Note that this is an implementation detail under the swap code, so from the perspective of callers, they don’t know that zswap exists. When a swap call happens, zswap runs the content of that page through zstd and writes the compressed data to a buffer.

Zswap has a dedicated memory allocator, called zsmalloc, which is in a way similar to a user-space allocator such as jemalloc. So after compressing the data, it requests that amount of memory from the allocator. If the request succeeds, it then copies the data there.

Zswap uses an index to map the swap identifier (device type, offset) to the virtual address where the compressed data ends up being stored, so if the time comes to swap in this page, zswap can look this data up, decompress it and go through the same process as if the data had been stored on disk.

The allocator zsmalloc can grow its memory dynamically, but it might determine at some point it’s too big, at which point it can fall back to regular swap.

Conclusion

This is one of the posts where I started with a very simple question but ended up in a rabbit hole, which caused the post to be much longer but also allowed me to learn a ton in the process.

Dealing with memory is particularly tricky because the metadata about memory is also stored in memory, so there’s a high degree of self-reference that can be hard to grok.

Related Posts

A lot of these memory-related flows are mentioned in Brendan Gregg’s books Systems Performance and BPF Performance Tools, which describe how to observe them. I recall some of the tools went over my head in terms of what they did, but having a deeper picture of these memory flows should help when I revisit them.

The post about ELF: Executable and Linkable Format discusses virtual memory from the process perspective.

In Local Inter-Process Communication, we briefly mentioned shared memory and memory-mapped files, both of which are relevant to this post but we didn’t cover them due to space. Shared memory is when physical addresses are mapped by multiple processes and poses extra challenges for memory reclamation.

Memory-mapped files are when disks have RAM semantics, so to make things more efficient, the kernel ends up keeping caches in memory which, as we saw, are the first to be reclaimed when under pressure.

Jacob Ziv and Abraham Lempel were both faculty at the Technion in Haifa, Israel, when in 1977 they published a paper called A Universal Algorithm for Sequential Data Compression in which they described a compression algorithm which, unlike Huffman coding, doesn’t rely on knowing the frequency of symbols upfront.

Ziv was an information theorist from the Electrical Engineering department while Lempel was part of the Computer Science department. They met due to their interest in lossless compression and Ziv said [1] that their skills complemented each other well:

I knew all about information theory and statistics, and Abraham was well-equipped in Boolean algebra and computer science.

This algorithm is now known as LZ77 (their last name initials + the year the paper was published). A popular algorithm called DEFLATE combines LZ77 with Huffman and is used by tons of software including gzip, git, png, etc.

In this post we’ll cover the LZ77 algorithm.

Intuition

Suppose we have a programming language with only two instructions: literal <symbol> and copy <rewind> <length>. The program starts with an empty output buffer. The instruction literal <symbol> adds <symbol> to the buffer. copy moves the cursor back <rewind> positions and copies the symbol under the cursor to the end of output <length> times.

Let’s do an example:

literal a; // a
literal b; // ab
literal c; // abc
copy 2 1;  // abcb
copy 2 2;  // abcbcb

The comments explain what’s going on. For the literals, we see it’s just adding symbols to the end of the output. For copy 2 1, we rewind the cursor to a[b]c and then copy it to the output abcb. We now rewind the cursor to ab[c]b and then copy the next two symbols to the end, abcbcb.

Here’s another example to highlight the overlap effect:

literal a; // a
literal b; // ab
literal c; // abc
copy 3 6;  // abcabcabc

The literals are the same as before. When we get to copy 3 6, and rewind 3 positions we have [a]bc. How is it possible to have <length> > <rewind>? The key thing is that we’re moving the cursor on a string that is growing at the same time, so when we add a to the output, we end up with a[b]ca, then ab[c]ab, then abc[a]bc, abca[b]ca, abcab[c]ab and finally abcabc[a]bc.

It would be a bit clearer if this language had, say, a rewind <rewind> command that moved the cursor <rewind> positions back and copy just copied the current symbol to the end and moved the cursor by one unit. Then our previous example would be:

literal a; // a[]
literal b; // ab[]
literal c; // abc[]
rewind 3;  // [a]bc
copy;      // a[b]ca
copy;      // ab[c]ab
copy;      // abc[a]bc
copy;      // abca[b]ca
copy;      // abcab[c]ab
copy;      // abcabc[a]bc

But this version is much more verbose. Even the original one we provided is very verbose. We can get away with having the literals as a single string, abc in our example, and copy now takes 3 arguments: lit_cnt, length and rewind. The only new one is lit_cnt which counts after how many literal instructions we inject this copy.

The previous example could be encoded as:

literal abc;
copy 3 6 3;

We’re basically saying the copy appears after all 3 literals which is not very instructive. Let’s consider a more complicated example with literal and copy interleaved:

literal a; // a
literal b; // ab
copy 2 2;  // abab
literal c; // ababc
copy 4 3;  // ababcbab
literal d; // ababcbabd

This could be represented more compactly as

literal abcd;
copy 2 2 2;  // inserted after 2 literals: a and b
copy 1 3 4;  // inserted after 1 literal: c

This is essentially the LZ77 encoding! The decoding part consists in “running” this program. The tricky part is the encoding, so let’s cover that.

Encoding

We’ll start with a greedy algorithm. We iterate over each character on the input. At a given position i, we try rewinding r positions and find the maximum prefix between a string starting at i and at i - r.

If the longest prefix found is at least 3 (this is because each copy requires 3 integers, so to be worth using it it must replace at least 3 characters) we emit a copy command for that prefix. Otherwise we emit the current position i as a literal.

In Python it could look like this:

def longest_prefix(s, i, j):
    """
    Finds the length of the longest common prefix
    between s[i:] and s[j:]. Assumes i >= j
    """
    p = 0
    while i + p < len(s) and s[i + p] == s[j + p]:
        p += 1
    return p

def best_rewind(s, i):
    """
    Finds the r such that the longest common prefix between
    s[i:] and s[i-r:] is the longest. Return r and the prefix
    length.
    """
    best_l = 0
    best_r = -1

    for r in range(1, i + 1):
        l = longest_prefix(s, i, i - r)

        if l > best_l:
            best_l = l
            best_r = r

    return best_r, best_l

def encode(input):
    i, prev_i = 0, 0
    lits, cps = '', []

    while i < len(input):

        best_r, best_l = best_rewind(input, i)

        if best_l >= 3:
            cps.append((i - prev_i, best_l, best_r))
            prev_i = i
            i += best_l

        else:
            lits += input[i]
            i += 1

    return lits, cps

The comments and code are easy to follow. The only observation is that when we append the triple (i - prev_i, best_l, best_r), we don’t store the current length of the literals but how many literals appeared between the current copy instruction and the previous one.

If $n$ is the length of the input, this algorithm is $O(n^3)$ in the worst case. We can bound this complexity by capping how far back we rewind in best_rewind() by changing:

for r in range(1, i + 1):
  ...

to

for r in range(1, min(i + 1, window)):
  ...

There’s another practical consideration in that we can’t afford to load the entire input into memory, so we can’t look forward or backward too far, so we can also limit how far longest_prefix() scans.

Decoding

Decoding is a bit simpler. As we discussed, it consists in running the program:

def decode(lits, cps):
    out = ''
    prev = 0
    for cnt, l, r in cps:
        out += lits[prev:prev + cnt]
        prev += cnt

        for _ in range(l):
            out += out[len(out) - r]

    out += lits[prev:]
    return out

Note that inside the innermost loop, len(out) changes as we append characters to it.

As for complexity, decoding is linear on the output, which we must emit anyway, so we can’t expect to do better than that.

Composition

As we mentioned at the start, LZ77 is typically not used on its own but rather combined with other encoders for further compression. We cover two popular composite encoders that use LZ77: DEFLATE and zstd.

At a very high level, DEFLATE applies LZ77 but uses the less compact form, the one that interleaves literals and the copy instructions. It counts the frequency of symbols and instructions on this output and then it applies Huffman coding.

The zstd algorithm applies LZ77 first using the compact form. It then either uses run-length encoding or Huffman coding for the literals and a variant of ANS called FSE (Finite State Entropy) for the copy instructions.

LZW

When I ran into LZ77, I could swear I had heard of Lempel Ziv before. Doing some searches on my email, I found that I had implemented the algorithm Lempel-Ziv-Welch or LZW in 2006, in assembly, for a class in college. I had no recollection of this implementation.

This is funny because when I first saw LZ77 encoding framed as a programming language with instructions literal and copy it reminded me of Assembly.

There’s an interesting story behind LZW. Lempel and Ziv improved upon LZ77 in 1978 and named it LZ78. In 1983, Terry Welch improved on LZ78 and called it LZW. Welch did so while working at Sperry Corporation, later Unisys, who patented the algorithm.

In 1987 CompuServe released the image format GIF, which used LZW. In 1994 Unisys started requiring software developers implementing LZW to pay license fees. Creating or distributing GIF images was not subject to this.

In 1993 Phil Katz invented the DEFLATE algorithm. As a response to the LZW controversy, the PNG image format was created in 1995 which used DEFLATE for compression. The patent for LZW expired in 2003 but due to the licensing requirements and DEFLATE’s better compression LZW isn’t widely adopted.

Conclusion

There are many tweaks and optimizations (e.g. not loading the entire input into memory) we can do but my main goal was to have a high level understanding of how LZ77 works! My main objective was to get a better sense of what goes into zstd.

References

• [1] IEEE Spectrum: Remembering Data Compression Pioneer Jacob Ziv.

Related Posts

It’s interesting that both LZW and Asymmetric Numeral Systems have some drama around patents. Encoders have an interesting trait that make them subject to this: once you use a given encoder, you are required to use the corresponding decoder. Another factor seems to be that encoding is low-level enough to be implemented in hardware, which has a stronger precedence of patenting vs. algorithms in general that cannot be patented because they’re abstract ideas.

If we look at the patent it’s actually describing hardware:

In the post Tree Ring Matching using the KMP Algorithm we’re trying to solve the problem of finding the largest overlap between a suffix and prefix of two strings. This is similar to finding the largest overlap between the prefixes in longest_prefix().

The LZ77 reminds me both of the KMP algorithm mentioned above and the Aho-Corasick where we have to efficiently find patterns in the input text.

CacheLib is an open-source C++ library from Meta for constructing in-process caches. At a high level, we can think of it as a hash table in memory with the option to spill to flash storage.

In this post we’ll study this library, mainly the in-memory portion based on the paper The CacheLib Caching Engine: Design and Experiences at Scale [1] and the code.

Motivation

The paper [1] states the problem CacheLib aims to solve: different systems reinvent the wheel when writing an in-process cache system. The argument for that is different systems require specialized cache solutions, so they can’t use off the shelf solutions.

CacheLib challenges that assumption by providing a flexible cache library that supports a multitude of production use cases including CDN, distributed application cache (e.g. memcached), general in-process cache, etc.

It also provides efficient multi-thread support and a secondary layer of flash storage.

Example

I always understand things better via examples, so let’s start with one, even though some parts need explaining later. The plan is to write a simple key to the cache, then query for that key and then another key that is not in there.

The API can feel a bit un-ergonomic because it’s meant to support multiple use cases, but we can define helpers that make it look more like a hash-table. First inserting an element:

bool put(
  Cache& cache,
  PoolId pid,
  const std::string& key,
  const std::string& value) {

  auto handle = cache.allocate(pid, key, value.size());
  if (!handle) {
    return false; // pool is full and nothing was evictable
  }
  std::memcpy(handle->getMemory(), value.data(), value.size());
  cache.insertOrReplace(handle);
  return true;
}

The API requires a pool, which we’ll cover later. When calling .allocate() it gets an ItemHandle if the library was able to find available memory. Then we copy the data from the value to the handle’s memory and finally signal to the handle to perform the write.

A natural question to ask is why the API is split into 3 pieces. I haven’t found an explicit motivation, but a legit one is that by giving you the memory buffer, you can construct the object directly there instead of copying. In our case we did end up copying anyway, but there might be cases where we don’t have to.

The get() is simpler:

std::optional<std::string> get(Cache& cache, const std::string& key) {
  auto handle = cache.find(key);
  if (!handle) {
    return std::nullopt;
  }
  auto *payload = reinterpret_cast<const char*>(handle->getMemory());
  return std::string(payload, handle->getSize());
}

The only observation is that we store raw bytes in the cache so we need to cast them back to the expected type. Here’s an example on how to write and read back from it:

#include <cachelib/allocator/CacheAllocator.h>

using Cache = facebook::cachelib::LruAllocator;
using facebook::cachelib::PoolId;

Cache::Config config;

config.setCacheName("example")
  .setCacheSize(256 * 1024 * 1024) // 256 MB of DRAM
  .validate();

Cache cache(config);

const auto pid =
  cache.addPool("default", cache.getCacheMemoryStats().ramCacheSize);

put(cache, pid, "hello", "world");

// retrieval
std::cout << get(cache, "my_key").value_or("<miss>") << "\n";
std::cout << get(cache, "other_key").value_or("<miss>") << "\n";

The size we set to setCacheSize() is a hard ceiling on how much the pools can consume. The pools themselves have a size, which in this case is the same as the cache size since we only have a single pool.

It’s worth mentioning that the size passed to the Cache::Config is the “physical” amount of memory we allow it to use, including for overhead such as metadata. The size we provide to the pool is the “logical” amount, how much it can use to store actual data.

To make the point clearer, we can’t do:

cache.addPool("default", 256 * 1024 * 1024);

We have to use cache.getCacheMemoryStats().ramCacheSize which is the effective amount of space available for use.

Concepts

Now that we have some idea on how to use CacheLib, let’s cover the major components and connect it to the example we saw.

Slab

The unit of memory used by CacheLib is called a slab, a chunk of contiguous 4MB memory. The slabs are 4MB-aligned too, meaning their starting address is divisible by 4MB. This allows us to optimize the logic to determine the slab address from a chunk address as we’ll see later (see Chunk).

Each slab carries 3 pieces of metadata: pool id (see Pool), class id (see Allocation Classes) and the chunk size (see Chunk).

Multiple items can be stored in a slab, but if its size is greater than 4MB, then an item is stored across multiple slabs connected as a linked list.

Slab Allocator

The slabs are stored in a contiguous array, which is initialized when creating the cache (based on setCacheSize()). The library uses mmap to get a contiguous virtual address space, but the mapping to physical space is done lazily by the OS.

This array is byte-aligned to 4MB meaning the address of each slab is divisible by 4MB or $2^{22}$. The maximum number of slabs the allocator supports is $2^{32} - 1$.

This array is owned by the slab allocator and it also maintains a list of free slabs.

Chunk

A chunk is a subdivision of a slab. We can think of each slab having an array of chunks of a fixed size, called allocation class and stored as the slab’s metadata.

Each chunk stores exactly one item. This means we need to have slabs with lots of different sizes to avoid internal fragmentation. Because if an item has size 75B and the smallest chunk size is 1K, all the remaining space is wasted. The minimum chunk size is 64B ($2^6$), so each slab can hold at most 65,536 ($2^{16}$) chunks.

In some situations we need to determine the slab address the chunk lives in with only the chunk address. Since we know slabs always have 4MB, first we can determine the slab index via:

// index of the slab in the slab array
idx  = (p - slabMemoryStart_) >> 22;
// address of the slab
slab = slabMemoryStart_ + (idx << 22);

But since we know it’s aligned at 4MB, we can simply zero the least significant 22 bits of the chunk address:

slab = p & ~((1 << 22) - 1);

Item

An item (CacheItem) is a class containing the key and value of the stored entry. It also stores metadata such as TTL. There’s a 1:1 mapping between a chunk and an item, so conceptually we can treat them as the same thing, as we’ll do throughout the post.

A chunk is a blob of memory, whereas the item is the data that lives in that location. A chunk has fixed size whereas the item size depends on what it’s storing. The chunk is not deleted when an item is evicted.

An item also represents a node in two linked lists, which we’ll call the index-linked-list and the eviction-linked-list. We’ll discuss them in Index and Eviction-Order, respectively. So while chunks have a fixed “physical” order inside a slab, the items have different logical orders for indexing and eviction purposes.

Allocation Classes

We can think of an allocation class as an object that manages all slabs with chunks of a given size. This is the ultimate object that determines which address to return back to the allocate() API.

It contains a few sets of information: a list of “active” slabs, a list of “free” slabs (empty slabs that it can use) and the current slab. All these are pointers to the global slab array. It also stores a stack of pointers to chunks that have been freed (freedAllocations).

An active slab never goes back to be a free slab. So as the program runs, and assuming the cache is fully utilized, we’ll get to an equilibrium state where there are no free slabs and each class is “right sized”. If item distribution changes though, there’s a background process to rebalance things (see Rebalancing).

Pool

A pool (class MemoryPool) is a collection of allocation classes. We can think of pools as strict partitions of the cache memory, because it has a specified size that is honored by the system. This is not true for allocation classes within a pool. As we’ll see in Rebalancing, the rebalancer can move slabs between allocation classes within a pool.

Pools are stored as an array by the pool manager (class MemoryPoolManager) and the pool id is just an index on that array. You can use multiple pools if you want the cache to be shared by different use cases with different quotas. At most 64 pools can be used, but in most cases a single pool is often enough.

The addPool() method hints at the responsibilities of the pool:

PoolId addPool(
    folly::StringPiece name,
    size_t size,
    const std::set<uint32_t>& allocSizes,
    const MMConfig& config = {},
    std::shared_ptr<RebalanceStrategy> rebalanceStrategy = nullptr,
    ...,
);

The allocSizes determines the allocation classes available to this pool, that is, the chunk sizes it supports. Note that it takes a set because it doesn’t make sense to have 2 classes with the same size. The MMConfig config is for the eviction policy (see Eviction-Order Container) and the others are related to rebalancing (see Rebalancing).

Index

The index is a hash table which conceptually maps a key to the chunk holding the value. The implementation is a fixed size array of 32-bit integers. The array size represents the number of buckets and it’s fixed throughout the cache lifetime. So when we insert a key into this array we hash the key to an integer, and take the modulo with the bucket count to find the index.

The 32-bit integer is used to store the location of a chunk. The 16 least significant bits identify the index of the slab in the global slab array while the 16 most significant bits identify the index of the chunk within that slab. With this addressing only $2^{16}$ or 65,536 slabs can be addressed, but we mentioned there can be up to $2^{32} - 1$ (see Slab). In this case a different address mechanism is used but we’ll not discuss it here.

To insert an entry in the index, the key is hashed and a bucket id is computed, and then the corresponding entry on the array points to the chunk address. If there was already a chunk at that bucket, the new chunk’s “next” pointer will now point to that previous chunk. In other words, on collisions we insert at the head of the linked list, like in a typical hash table.

Eviction-Order Container

Conceptually an eviction-order container is a doubly linked list containing the elements of a given pool-allocation class pair. Each such container implements different eviction policies, but let’s assume it’s LRU (least recently used) for simplicity.

At any given time, the order in this list represents the eviction order, such that the last element on the list will be evicted first when needed. So assuming LRU, whenever an item is inserted in the cache, it’s inserted at the head. The container has a pointer to both the head and tail of the list.

The problem with eviction policies is that we often need to update the order when we read from the cache. Suppose we read a given chunk belonging to a linked list. We need to move it to the head of the list. This can be done in $O(1)$ because we get access to the chunk via the index and the “prev” and “next” pointers are part of the chunk. So we just do some pointer stitching to move it to the head.

Data Flows

Let’s now cover a few data flows, the main one being insertion and see how the components fit together.

Insertion

When we call cache.allocate(pid, key, size) like in the example, it will internally ask the pool corresponding to pid for a chunk of at least size amount of memory. The pool will look at the different size classes it has and will ask the one with the smallest size that is greater or equal to size.

Within the class, the request for a new chunk is: if freedAllocations is not empty it returns the top of the stack. Otherwise, it gets a new chunk on the current slab. If the current slab is full, it chooses a free slab and makes it current. If no free slabs remain, the pool will try to give it more slabs.

The pool has its own set of free slabs. If it has it available, it gives it to the allocation class. If not and it’s within its budget, it requests to the slab allocator which has a global set of free slabs.

If no such slab can be found, then the cache might try to evict items from the allocation class. In Eviction-Order Container we mentioned how chunks are kept in “eviction order” so it will give that chunk to the requester. It’s not always possible to evict a chunk so it may still fail to return one.

It will not insert the item in the chunk until insertOrReplace() is called, so it will not update the index-linked-list nor the eviction-linked-list until this method is called. This means this item is not “visible” to other threads until then, but its chunk is reserved.

Existing Key. One interesting difference in behavior from a hash table is that allocate() does not check if the key already exists in the cache. It first acquires a chunk and then when calling insertOrReplace() the old item is removed.

Eviction

Let’s focus on a flow mentioned during insertion: eviction. As we’ve seen the most common way to evict an item is when we run out of available space. Before we return it to the requester, we need to remove the chunk from the eviction order container’s doubly linked list and from the linked list on the allocator class.

An entry can be explicitly removed via remove() or via insertOrReplace(): recall that it allocates a new chunk and then removes the old one, it doesn’t update the chunk in-place. In these cases, the chunk is also added to the freedAllocations of its allocation class. It might need to update the index too (if this was the first chunk on the list).

There’s also a background thread, the reaper, which runs every 5 seconds by default. It traverses the array of slabs and within each slab the array of chunks. Within each chunk it will check the item’s TTL and if it’s expired it gets evicted.

Lookup

A lookup for a key consists of hashing the key, finding the right entry in the index and then traversing the corresponding linked list formed by the items, until we find the one matching the key. A handle is then returned to that item from which we can read the data.

As we discussed in Eviction-Order Container, the item might need to be moved depending on the eviction policy. For policies such as LRU it can be done in $O(1)$.

Rebalancing

Rebalancing is done by a background thread that runs every second. For each pool, it transfers slabs from the most empty class to the most full one.

If the source class has slabs in its free slabs list, the move is straightforward. If any of the active slabs is actually empty, it’s chosen. If it has some items in it, they’re std::moved to a different slab (defragmented) and this slab is chosen.

Concurrency

In describing the data flows above we glossed over the intricacies of concurrency. CacheLib is designed to support multiple threads at a time so some designs that may seem arbitrarily complex might make sense in light of thread safety.

Index and Eviction-Container Locks

The first layer of lock is on the index table. In theory we could have one lock per row, but if we use a mutex, it’s 4 bytes per row, which would double the memory size of the index table. On the other hand, if we locked the entire table with a single lock we’d run into massive contention. A middle ground is for a lock to be responsible for a subset of rows. There’s contention only for rows on the same set, but we use fewer locks.

During a lookup, it acquires a read lock on the index to find the right chunk to retrieve. Note that this conceptually gives a read lock to the index-linked-list associated with that bucket. Implementation wise, the “next” pointer for the index-linked-list can only be used after acquiring the appropriate lock.

To “move” the item within an eviction-order container, a write lock is acquired, which analogously gives write access to the eviction-linked-list, and that is implemented by requiring a lock to read/write the “prev” / “next” pointer for that list.

For the put() it needs to acquire a write lock on the index in addition to write locks on the eviction-order container. There are several other locks involved, especially when a new slab is needed, but the linked list ones are the more interesting.

Item

The ItemHandle returned by cache.allocate() has a reference count. If the reference count is non-zero, then the cache won’t evict the corresponding item.

Recall from Data Flows > Insertion that when we do cache.allocate(pid, key, size) we don’t update the linked lists, so effectively the ItemHandle is not visible to other threads. So doing

std::memcpy(handle->getMemory(), value.data(), value.size());

in put() is thread safe. Note that it’s possible multiple threads call put() on the same key, and the last one to call insertOrReplace(), which is thread safe, will win.

Flash

All we’ve discussed so far is about the structure in memory. CacheLib has a second layer of caching backed by Flash. We won’t cover too many details in this post though.

When an eviction happens in memory, there’s an option to save the data in flash. There are some tricks described in [1] to avoid writing to flash excessively due to limited write lifetime.

Once the entry is written to flash it’s also written to an in-memory per-bucket Bloom filter. During lookup, when the in-memory lookups fail, it will check if the entry is on the bloom filter.

If yes, there’s a chance it exists in flash, so it tries to read the entry from flash. If a match is found, the entry that was in flash is moved to memory again.

Conclusion

I’m having a lot of fun studying these data structures such as CacheLib and F14Map. They have lots of interesting tricks to make it efficient and perform well under concurrent use. I especially liked the fact that an item is a node in multiple linked lists.

These complex data structures also make me think of the joke that tech interviews ask interviewees how to reverse a linked list but that no one uses that in their job. To be fair this is pretty rare or unheard of for many programmers but I found the deeper one goes on the stack the more likely they’ll need Computer Science knowledge.

I was surprised that the paper focuses so much on unification and use cases instead of the technical details (the exception seems to be about writing to flash). Maybe because these are standard tricks in cache implementation, but they felt all novel to me.

References

• [1] The CacheLib Caching Engine: Design and Experiences at Scale
• [2] Github: CacheLib

Related Posts

There are a lot of similarities in how systems memory allocators manage memory. One of them is hierarchy: we see them in arenas in Jemalloc and pools in Velox: Memory, the latter having an entire tree of hierarchy. Velox also has disk spill capabilities, much like CacheLib can “spill” to flash.

At a very high level, an in-process cache is a hash table, but it has substantial differences from a general purpose hash map such as Folly F14 Map. It typically has a fixed size and it evicts entries automatically, which leads to very different design decisions.

Memory allocators also have to worry about fragmentation, which can lead to inefficient use of space. This is also a concern of the Buddy Memory Allocation.

In terms of similar use cases, the CPU also has a cache, as we’ve studied in CPU Cache. We also discussed CDNs before in Content Delivery Network.

Finally, speaking of nodes that belong to multiple linked lists, we talked about Dancing Links in The Algorithm X and the Dancing Links which is a way to represent sparse matrices.

Jarosław Duda is a Polish professor at the Jagiellonian University in Kraków. He developed a family of entropy coding methods called asymmetric numeral systems (ANS), mainly used in data compression. He wanted these to remain patent-free but has had mixed success.

Google worked with Duda around 2014 in a paper Mixed boolean-token ANS coefficient coding and tried to patent a coder for video, but Duda pushed back and Google abandoned the attempt. In 2019 Microsoft was able to patent a variant and since then other patents have been granted internationally.

A variant of ANS, known as FSE, is used by Meta’s compression library called zstd. I wanted to learn more about it and decided to study this algorithm first.

Setup

We’ll develop the Python code as we follow along, both to make the explanations more precise and to define common boilerplate gradually so each code snippet stays shorter. Since we’ll be developing and comparing encoders, we define an abstract class:

class Encoder(ABC):
    @abstractmethod
    def desc(self) -> str:
        "Human-friendly description of this encoder"

    @abstractmethod
    def encode(self, input: str) -> bytes:
        "Gets an input, encodes to bytes"

    @abstractmethod
    def decode(self, encoded: bytes) -> str:
        "Decodes bytes to the string"

A dummy encoder is simply converting strings to bytes using UTF-8. Not very useful but it gives us a working implementation:

class UTFEncoder(Encoder):
    def desc(self) -> str:
        return "UTF-8"

    def encode(self, input: str) -> bytes:
        return input.encode('utf-8')

    def decode(self, encoded: bytes) -> str:
        return encoded.decode("utf-8")

A simple way to evaluate the encoders is to encode an input, decode it and verify it matches the original input since we’re assuming lossless encoders. We can also compute the compression rate as (1 - encoded/decoded) and the time it takes to encode and decode.

encoders = [
    UTFEncoder(),
]

decoded_sz = len(data)
for encoder in encoders:
    start = time.perf_counter()

    encoded = encoder.encode(data)
    encoded_sz = len(encoded)
    rate = (1 - encoded_sz/decoded_sz) * 100

    output = encoder.decode(encoded)
    assert data == output

    elapsed = time.perf_counter() - start

    rows.append([
        encoder.desc(),
        elapsed,
        rate,
    ])

print(tabulate(
    rows,
    headers=["Encoder", "Time (s)", "Compression Rate"],
    tablefmt="rounded_grid",
    floatfmt=".3f",
))

Huffman Coding

We wrote about Huffman coding before. The idea is: given a probability distribution of symbols, we construct the Huffman tree which is a binary tree where symbols are at the leaves. Then we build a map from each symbol to a binary code representing the path to the corresponding leaf (0 is left, 1 is right).

We can then encode an input string character by character as a byte string. This encoding has the property of being prefix-free, which allows us to decode the byte string greedily by using the 0s and 1s to traverse the Huffman tree.

The Huffman encoding minimizes this function:

where $p_i$ is the probability of symbol $i$ and $l_i$ is the length of the path from the root to the leaf containing $i$. We then talked about entropy which is defined as:

and that while the Huffman encoding doesn’t minimize it, it’s within $H + 1$.

Intuitively the gap is that $l_i$ is an integer, while $\log(p_i)$ is a real value. This is where Asymmetric Numeral Systems come in. Before we go there, let’s cover a more familiar concept but with an unfamiliar name: symmetric numeral systems.

We won’t cover the code for Huffman again, but we can assume a class HuffmanEncoder exists extending Encoder.

Symmetric Numeral Systems

Consider how we convert a binary string $b_n b_{n-1} b_{n-2} \cdots b_0$ (most significant bit first) to decimal. In Python it can look like:

def btoi(bits):
    r = 0
    for b in bits:
        r = r * 2 + int(b)
    return r

Note how we use one bit per symbol 0 or 1 and note how the integer roughly doubles each time. Another way to implement this is via:

def encode_bit(x, b):
    match b:
        case '0':
            return 2 * x
        case '1':
            return 2 * x + 1

def btoi(bits):
    r = 0
    for b in bits:
        r = encode_bit(r, b)
    return r

One way to interpret encode_bit() is as an explicit map where 0 gets the even integers and 1 gets the odd ones. Then we create a mapping from the natural numbers to their respective subset.

    0  1  2  3  4  5 ...
0:  0  2  4  6  8 10 ...
1:  1  3  5  7  9 11 ...

Another way to see this is as a partition of the natural numbers into two subsets and the association of one to 0 and another to 1. As we know, this function is reversible so we can convert an integer to a string of bits (if we know how many bits we started with).

A different way to see this is as a reversible operation is to look at which partition of the naturals the integer we’re decoding is at. If it’s an even number, it means it was last mapped by a 0, whereas if it’s odd, it was mapped by a 1. This interpretation will be more useful soon.

We can generalize this encoding for an arbitrary base, and thus we can also encode an input that used any alphabet. For example, if we have 3 symbols and the input ACBAABCBA, we can encode it as an integer by converting this as if it was a string of digits in base 3.

Asymmetric Numeral Systems

Suppose now that 0s are 3x more likely than 1s to appear in the input. Intuitively we want to use fewer bits to encode 0 than 1. We can now distribute the image non-uniformly: for every 4 numbers we assign the first 3 to 0 and the fourth to 1:

    0  1  2  3  4  5 ...
0:  0  1  2  4  5  6 ...
1:  3  7 11 15 19 23 ...

In the original example, we said each input maps to roughly 2x its value. Now for 1 the input grows at 4x rate, but for 0 it’s more like 4/3x, which is closer to their distribution!

We can represent this distribution in encode_bit() via:

def encode_bit(x, b):
    match b:
        case '0':
            return 4 * (x // 3) + (x % 3)
        case '1':
            return 4 * x + 3

This numeral system, in which different digits have different weights, is called Asymmetric Numeral Systems or ANS. As in the symmetric case, we can think of a general base $n$ system in which digit/symbol $i$ occurs with probability $p_i$.

We find a suitable denominator and then the rational that is the closest approximation to the probability. For example, if we have 3 symbols (A, B, C) with probabilities $0.5$, $0.3$ and $0.2$ respectively and we choose a denominator $M = 16$, we can have $A = 8 / 16$, $B = 5 / 16$ and $C = 3 / 16$.

We can then distribute a chunk of 16 integers between them: the first 8 goes to A, the next 5 to B and the last 3 to C:

0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15
A  A  A  A  A  A  A  A  B  B  B  B  B  C  C  C

Let $M$ be the chunk size, $f_i$ how many slots of $M$ we assign to symbol $i$, and $o_i$ the offset in the chunk where $i$ starts. So in the example above we have: $f_A = 8, f_B = 5, f_C = 3$, $o_A = 0, o_B = 8, o_C = 13$.

To find to which integer an input $x$ maps for a symbol $i$, we can first find the right chunk by doing $\lfloor{x / f_i\rfloor}$. So for $A$, 0-7 will map to the 1st chunk, 8-15 to the 2nd, etc. Then we compute the offset of that chunk by multiplying by $M$. Now we add the offset of $i$ within that chunk, $o_i$ and finally the offset of the $k$-th copy of $i$ in that segment, $k = x \pmod {f_i}$. The formula is a bit convoluted but you can try with a few examples to convince yourself.

Now to the implementation. We’ll first define the class SimpleANSEncoder extending Encoder, which computes $f$ (as freq) and $o$ (as offsets):

FrequencyTable = Sequence[tuple[str, int]]

class SimpleANSEncoder(Encoder):

    def __init__(self, raw_freq: FrequencyTable) -> None:
        self.M = 2**12

        # least frequent first
        sorted_freq = sorted(raw_freq, key=lambda item: item[1])
        total = sum(cnt for _, cnt in sorted_freq)

        freq_norm = [
            (s, max(int(cnt/total * self.M), 1)) for s, cnt in sorted_freq
        ]

        self.offsets = {}
        acc = 0
        for s, f in freq_norm:
            self.offsets[s] = acc
            acc += f

        # tweak last symbol freq so it adds up to M
        last = freq_norm[-1]
        freq_norm[-1] = (last[0], last[1] + self.M - acc)

        self.freq = dict(freq_norm)

The code is a bit complicated because we want to make sure that sum(prob) == M, while we’re rounding floats to integers. A trick is to have the most frequent symbol have M - the sum of the probabilities of the least frequent items.

For incorporating a symbol into the integer x we follow the formula above:

def encode_symbol(self, x: int, symbol: str) -> int:
    f = self.freq[symbol]
    off = self.offsets[symbol]
    return self.M * (x // f) + (x % f) + off

The encode() consists in building the giant integer and encoding it as a byte string. We can use BytesIO to read/write to this byte string:

def encode_int(cursor: BytesIO, x: int) -> bytes:
    cursor.write(x.to_bytes(4, byteorder="big"))

def encode(self, input: str) -> bytes:
    r = 0
    for c in input:
        r = self.encode_symbol(r, c)

    cursor = BytesIO()
    encode_int(cursor, len(input))
    encode_bigint(cursor, r)
    return cursor.getvalue()

We also include the length of the original text because otherwise the decoder doesn’t know how many symbols to generate back. For example, in the binary to decimal encoding, leading zeros would be lost during encoding, so 0001 gets encoded as 1 and decoded as 1.

The decoding is only a bit more complicated. Finding the symbol is easy: we just need to mod $M$, since this gives the offset within a chunk and then we can do a binary search on off to find to which offset it belongs. Alternatively we store an array of size $M$ with the symbols repeated, which we can compute upfront, in __init__(). In our example:

class SimpleANSEncoder(Encoder):
    def __init__(self, raw_freq: FrequencyTable) -> None:
        # ...
        self.lookup = "".join(s * self.freq[s] for s, _ in sorted_freq)

    def get_symbol(self, x):
        return self.lookup[x % self.M]

The decode_symbol() is roughly the inverse of encode_symbol once we have the symbol:

def decode_symbol(self, x: int) -> tuple[str, int]:
    symbol = self.get_symbol(x)
    p = self.freq[symbol]
    off = self.offsets[symbol]
    return symbol, p * (x // self.M) + (x % self.M) - off

And so is the decode():

def decode_int(cursor: BytesIO) -> int:
    return int.from_bytes(cursor.read(4), byteorder="big")

def decode(self, encoded: bytes) -> str:
    cursor = BytesIO(encoded)
    sz = decode_int(cursor)
    value = decode_bigint(cursor)

    output = ''
    for _ in range(sz):
        symbol, value = self.decode_symbol(value)
        output += symbol
    return output[::-1]

Efficiency

The rational approximation probability of a symbol is $q_i = f_i / M$. The rough growth of x when going through encode() is roughly dominated by $M / f_i$, so $1 / q_i$. The approximate number of bits in this factor is $\log_2 1/q_i = - \log_2 q_i$. So the expected cost per symbol is a weighted sum:

Where $L$ is the cost of a random variable representing a symbol. This expected value is also equal to the cross-entropy which indicates how many bits on average we need if the symbols come from a distribution $P$, but we encode it using a distribution $Q$ and is denoted by $H(P, Q)$.

We can compare this with the true entropy:

which we’ve seen in Huffman coding is the theoretical optimal. So the closer we can approximate the rational $q_i$ to $p_i$ the better, so in theory we could use a gigantic value for $M$ but in practice working with large values is prohibitive.

Speaking of large integers, encoding an input as an integer has the major downside of the integer growing exponentially with the size of the input.

Renormalization

One way to avoid working with arbitrarily large integers is to move the lowest bits to a bit stream whenever it grows too big. It will be part of the encoded output but is not going to participate in future multiplications. The idea is to have the invariant $M \le x \lt 2M$. If multiplying $x$ by a factor would tip it over $2M$ we first reduce it.

We define a class RANSEncoder extending from SimpleANSEncoder since they share a lot of the methods. For encoding, encode_symbol() is the same, but now we need to emit bits periodically. We use bitarray for that:

def encode_bitarray(cursor: BytesIO, x: bitarray) -> None:
    encode_int(cursor, len(x))
    cursor.write(x.tobytes())

class RANSEncoder(SimpleANSEncoder):

    # ...

    def encode(self, input: str) -> bytes:
        stream = bitarray()
        x = self.M
        for b in input:
            while (y := self.encode_symbol(x, b)) >= 2 * self.M:
                stream.append(x & 1)
                x >>= 1
            x = y

        cursor = BytesIO()
        encode_int(cursor, len(input))
        encode_bitarray(cursor, stream)
        encode_bigint(cursor, x)
        return cursor.getvalue()

Note how we start $x$ as $M$ as opposed to $0$ as we did before, so that we maintain the invariant $M \le x \lt 2M$ inside the loop.

For decoding, we can also maintain the invariant and consume bits from the stream to restore it:

def decode_bitarray(cursor: BytesIO):
    x = bitarray()
    bit_cnt = decode_int(cursor.read(4))
    byte_cnt = (bit_cnt + 7) // 8
    x.frombytes(cursor.read(byte_cnt))
    x = x[:bit_cnt]  # remove padding
    return x

class RANSEncoder(SimpleANSEncoder):

    # ...

    def decode(self, encoded: bytes) -> str:
        cursor = BytesIO(encoded)
        input_size = decode_int(cursor.read(4))
        stream = decode_bitarray(cursor)

        out = ''
        x = decode_int(cursor.read())
        for _ in range(input_size):
            s, x = self.decode_symbol(x)
            while x < self.M:
                bit = stream.pop()
                x = (x << 1) | bit
            out += s

        return out[::-1]

Not only does this allow us to operate with a small integer range, the set of values we pass for x to encode_symbol() is bounded between [M, 2M[, which is very useful as we’ll see next.

Tabular ANS

Recall that in our 3-symbol example with distinct weights we laid them out like this:

0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15
A  A  A  A  A  A  A  A  B  B  B  B  B  C  C  C

First the A’s, then the B’s and then the C’s. The problem with this approach is that C always gets penalized. For example, if the input is, say 1, and the symbol is C, it will be mapped to 14, while B gets mapped to 9 and A gets mapped to 1. So C grows a lot faster.

It will be more balanced if we interleave the symbols more evenly within one batch, for example:

0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15
A  B  A  C  A  B  A  B  A  C  A  B  A  B  A  C

Now for input 1, A maps to 2, B to 5 and C to 9. The problem with this arbitrary order is that it becomes much harder to devise a formula that maps integers to their corresponding output. That’s where the renormalization helps: if we have a fixed range of values of input ([M, 2M[), we can write a table with M entries that map them exactly to the order above! This variant is called the Tabular ANS.

Let’s create the TRANSEncoder also extending from SimpleANSEncoder. The main difference on the initialization is that we’ll “shuffle” the lookup array and then pre-compute the encode and decode table.

class TRANSEncoder(SimpleANSEncoder):
    def __init__(self, raw_freq: FrequencyTable) -> None:
        super().__init__(raw_freq)

        step = (self.M >> 1) + (self.M >> 3) + 3
        lookup: list[str] = [''] * self.M

        pos = 0
        for symbol, cnt in self.freq.items():
            for _ in range(cnt):
                lookup[pos] = symbol
                pos = (pos + step) % self.M

        self.lookup = lookup
        self.precompute_encode()
        self.precompute_decode()

The method precompute_encode() is analogous to RANSEncoder’s encode() but it uses the positions computed above. It returns a table that for each M <= x < 2*M and symbol s, stores a pair (x, bits), where x is the new value of x after incorporating s and the bits that will go into the stream.

def precompute_encode(self) -> None:
    positions = defaultdict(list)
    for p, s in enumerate(self.lookup):
        positions[s].append(p)

    def encode_symbol(x: int, s: str) -> int:
        slots = positions[s]
        f = len(slots)
        return self.M * (x // f) + slots[x % f]

    self.encode_table: list[dict[str, EncodeEntry]] = [
        {} for _ in range(self.M)
    ]
    for x0 in range(self.M, 2 * self.M):
        for s in self.freq:
            x = x0

            slots = positions[s]
            f = len(slots)
            bits = bitarray()
            while (y := encode_symbol(x, s)) >= 2 * self.M:
                bits.append(x & 1)
                x >>= 1

            self.encode_table[x0 - self.M][s] = (
                y,
                bits,
            )

With this table, encoding becomes very simple (omitting the writes):

def encode(self, input: str) -> bytes:
    stream = bitarray()
    x = self.M
    for b in input:
        x, bits = self.encode_table[x - self.M][b]
        stream += bits
    # ...

The decode table only depends on M <= x < 2*M and it maps to a pair: the previous state and how many bits it needs to consume from the stream to grow x back to the range.

def precompute_decode(self) -> None:
    cnt = defaultdict(int)
    decode_table = []
    for s in self.lookup:
        rank = cnt[s]
        cnt[s] += 1

        prev_x = self.freq[s] + rank
        bit_cnt = 0
        while (prev_x << bit_cnt) < self.M:
            bit_cnt += 1

        prev_x <<= bit_cnt

        decode_table.append((
            prev_x,
            bit_cnt
        ))
    self.decode_table  = decode_table

Decoding becomes:

def decode(self, encoded: bytes) -> str:
    cursor = BytesIO(encoded)
    input_size = decode_int(cursor)
    stream = decode_bitarray(cursor)

    out = ''
    x = decode_bigint(cursor)
    for _ in range(input_size):
        off = x - self.M
        s = self.get_symbol(off)
        base, bit_cnt = self.decode_table[off]

        suffix = 0
        for _ in range(bit_cnt):
            suffix = (suffix << 1) | stream.pop()

        x = base | suffix
        out += s

    return out[::-1]

Experiments

I used Carroll’s Alice in Wonderland as a training dataset to estimate the character frequencies, then encoded Shakespeare’s As You Like It (125k characters):

As expected, the no-op UTF-8 is pretty efficient! But it doesn’t compress anything. Huffman is very fast and is only slightly worse compression than ANS. Working with gigantic integers is very slow.

tANS was slower than rANS mostly because the overhead of constructing the tables wasn’t enough to pay off for 125k characters. To verify that, I also tried encoding a bigger text ($10^8$ characters):

Here the overhead paid off for tANS. It also achieved a better compression rate with the more spread out symbols in the chunk.

Conclusion

I found it hard to grok ANS at first until realizing it is, in a way, a generalization of the binary to decimal algorithm!

I found it difficult to implement these in Python, and got a lot of “off-by-one” errors. Luckily, testing a lossless encoder is relatively easy, we just need to test the result is the same as the input.

Related Posts

In The Cardinality of Complex Numbers we talk about bijective mapping between $\mathbb{C}$ and $\mathbb{R}$, which one can see as a lossless encoding. This also vaguely reminds me of the Gödel number system which maps mathematical expressions to natural numbers!

One of the most surprising coincidences is that in the recent post about Folly F14 Map we discussed the idea of “shuffling” elements across chunks much in the same way TRANSEncoder shuffles the symbols over the array. Both choose the size to be a power of 2 and the stride or step to be an odd number.

Another connection is that Folly F14 Map is a hash table and the first hashing algorithm I learned and possibly the only one I ever implemented is related to Symmetric Numeral System: we treat the key string as a number in base, say 128 (for ASCII), and then convert it to an decimal modulo a prime number.

Like Huffman, ANS is a lossless compression, but in many cases a lossy compression is a tradeoff we can take, such as in Linear Predictive Coding and T-Digest.

Finally, a weaker connection is with Numerical Representations as inspiration for Data Structures, because it’s also based on number systems.

Otakar Borůvka was a Czech mathematician who is best known for his work in graph theory. Once, his friend Jindřich Saxel, an employee of the West Moravian Power Company, asked him for help optimizing electric distribution networks.

Borůvka modeled the problem as the minimum spanning tree problem and then came up with the first known algorithm to solve it, now known as the Borůvka algorithm.

In this post we’ll explore the Borůvka algorithm combined with KD-trees to solve the Euclidean Minimum Spanning Tree problem more efficiently.

Problem

The Euclidean Minimum Spanning Tree problem (EMST) consists in finding the minimum spanning tree (MST) of a set of $N$ points in the Euclidean space.

Algorithms such as Kruskal can find the MST of a graph $G(V, E)$ in $O(\abs{E} \log \abs{E})$. However in EMST the edges are implicit. The graph is a complete one where the weight of the edge between two points is their Euclidean distance, so a naive implementation consisting in computing the edges explicitly leads to a $O(N^2 \log N)$ algorithm.

Solution

In this post we’ll explore a more efficient implementation using the Borůvka algorithm with KD-trees which runs closer to $O(N \log^2 N)$ on average, even though the worst case is $O(N^2 \log N)$.

We’ll start by describing the Borůvka Algorithm to compute the MST. Note that there’s nothing special about it with respect to the Euclidean variant. It’s just that it operates in batches and we can leverage that to search on the KD-trees more efficiently.

Borůvka Algorithm

Interestingly enough, the problem that motivated the Borůvka algorithm can be modeled as a Euclidean graph, where vertices are cities and edges between them have costs corresponding to the Euclidean distance. However his algorithm works on general graphs and doesn’t exploit the Euclidean properties.

Like Kruskal and Prim, the Borůvka algorithm is greedy, at each step taking the optimal decision and that leads to the optimal solution. In high level terms:

At start, each vertex belongs to its own component. Then we loop until one component remains:

• For each component, select the cheapest edge to another component. Collect all these edges first.
• For each edge:
  • If they connect two components, merge them into one, add the edge to the solution.

This algorithm feels like a fusion between Kruskal (add cheapest edges to the MST) and Prim (grow a component). I guess it’s not as popular as either because the implementation is less simple. The steps are very suggestive of the union-find data structure and indeed it’s used in the implementation.

First, let’s prove the correctness of this algorithm: the idea is that at any point in time, the solution contains a subset of a MST (Lemma 1). This then implies the edges in the solution form a forest (i.e. each component is a tree). Since only one component is left at the end, it must be a tree and hence a MST.

Lemma 1. At each step of the iteration of the Borůvka algorithm the set of selected edges is a subset of the edges of some spanning tree of minimum cost.

The simplest way to implement this algorithm is to keep a list of all candidate edges. Then at the beginning of the iteration, keep only edges that are across components (using union-find) and then for each component, track the cheapest edge. We can determine whether two vertices are in the same component of a union find in $O(1)$ amortized, and also merge two components in $O(1)$. This makes each iteration $O(\abs{E})$.

At each iteration, each component will be merged into some other component, so the number of components will at least halve. Thus there are only $O(\log \abs{V})$ iterations leading to a $O(\abs{E} \log \abs{V})$ algorithm.

Implementation

As discussed above, the algorithm uses the union find data structure, which we’ll implement using UnionFind. The implementation details are collapsed below. We just need to understand the methods: .find(x) -> int which returns the component of a given member x, and .union(x, y) which merges the components of x and y.

class UnionFind:
    def __init__(self, n):
        self.p = list(range(n))
        self.s = [1]*n
        self.sz = n

    def find(self, x):
        p = self.p
        if p[x] != x:
            p[x] = self.find(p[x])
        return p[x]

    def union(self, x, y):
        px, py = self.find(x), self.find(y)
        if px == py:
            return

        s, p = self.s, self.p
        if s[px] < s[py]:
            px, py = py, px

        s[px] += s[py]
        p[py] = p[px]

        self.sz -= 1
        return s[px]

    def size(self):
        return self.sz

We also define a few utility classes to help with syntax, but their meaning is intuitive, for example Point and BoundingBox. A bounding box is the smallest axis-aligned box containing a set of points. We represent it by its “lower” and “upper” corners.

def norm_sq(p) -> float:
    return sum(x*x for x in p)

@dataclass(frozen=True)
class Point:
    data: tuple[int, ...]

    def __getitem__(self, dim: Dimension):
        return self.data[dim.v]

    def __len__(self) -> int:
        return len(self.data)

    def __abs__(self) -> float:
        return sqrt(norm_sq(self))

    def __sub__(self, other: Point):
        sub = tuple(a-b for a,b in zip(self.data, other.data))
        return Point(sub)

    def __iter__(self):
          return iter(self.data)

@dataclass
class BoundingBox:
    lo: Point
    hi: Point

    @staticmethod
    def for_point(p):
        return BoundingBox(p, p)

    def union(self, bb):
        if bb is None:
            return self

        self.lo =  Point(
            tuple(min(px, qx) for (px, qx) in zip(self.lo, bb.lo))
        )
        self.hi =  Point(
            tuple(max(px, qx) for (px, qx) in zip(self.hi, bb.hi))
        )

    def size(self):
        'Square length of the diagonal'
        return norm_sq(self.lo - self.hi)

We can start defining an interface for any algorithms computing the Euclidean minimum spanning tree from a set of points:

class EMST(ABC):
    @abstractmethod
    def desc(self) -> str:
        "Describe the algorithm"

    @abstractmethod
    def get_mst_edges(self, points: list[Point]) -> list[Edge]:
        pass

And then define the specialization for Borůvka:

class Boruvka(EMST):
    def setup(self, points: list[Point]):
        pass

    @abstractmethod
    def fill_best_edge_by_component():
        pass

    def get_mst_edges(self, points: list[Point]) -> list[Edge]:
        n = len(points)
        uf = UnionFind(n)
        self._uf = uf

        self.setup(points)

        sol = []
        while uf.size() > 1:
            self._best_edges = {}
            self.fill_best_edge_by_component()
            for e in self._best_edges.values():
                uc, vc, _ = e
                if self.merge_components(uc, vc):
                    sol.append(e)
        return sol

The core of this algorithm is fill_best_edge_by_component(), i.e. for each component in the union-find, determine the shortest edge from it to some other component. From there we determine which edges to keep and which components to merge.

We also add some helper functions that are used by different Borůvka implementations, mostly to avoid repeated code when dealing with the union find structure such as merge_components() and process_edge() (see if it’s worth adding an edge to the best edge by component).

class Boruvka(EMST):
    ...
    # Helper methods

    def merge_components(self, p, q) -> bool:
        'Combine components of p and q into one'
        if self._uf.find(p) == self._uf.find(q):
            return False
        self._uf.union(p, q)
        return True

    def get_component(self, p):
        return self._uf.find(p)

    def get_best_dist(self, p):
        pc = self.get_component(p)
        if pc not in self._best_edges:
            return inf
        return self._best_edges[pc][2]

    def process_edge_for_component(self, c, e):
        _, _, w = e
        if self.get_best_dist(c) > w:
            self._best_edges[c] = e

    def process_edge(self, e):
        u, v, w = e
        uc, vc = self.get_component(u), self.get_component(v)
        if uc == vc:
            return

        self.process_edge_for_component(uc, e)
        self.process_edge_for_component(vc, e)

Now we are ready for our first implementation of the “naïve” Borůvka, one where we compute the set of edges explicitly. With the helpers we defined, the code is straightforward: to compute the best edge set, we just process each of the $O(N^2)$ edges one at a time:

class NaiveBoruvka(Boruvka):
    def desc(self) -> str:
        return "Naive Boruvka"

    def setup(self, points: list[Point]):
        n = len(points)
        edges = []
        for i in range(n):
            for j in range(i + 1, n):
                edges.append((i, j, dist(points[i], points[j])))
        self._edges = edges


    def fill_best_edge_by_component(self):
        for e in self._edges:
            self.process_edge(e)

KD-Tree

In KD-Tree we learned that this binary structure allows us to search for the nearest point by splitting clusters of points into different subtrees.

It works well in practice, closer to $O(\log N)$ search, even though in the worst case it’s $O(N)$. We can use it to find the closest point $p$ to each point $q$ so that the edges $(p, q)$ are the only ones that need to be accounted for. Conceptually we could implement fill_best_edge_by_component() as:

class KDTreeBoruvka(Boruvka):
  ...
  def fill_best_edge_by_component(self):
      for q in points:
          p = self.get_closest_point(kd, q)
          self.process_edge((q, p, dist(p, q)))

Since on average we expect each search to be $O(\log N)$, fill_best_edge_by_component() is $O(N \log N)$ and the Borůvka becomes $O(N \log^2 N)$.

We can’t use a generic search for kd-tree because now we need to account for the components of the points: Points on the same component are not candidates for closest neighbors. The difference with query_kd_tree() in the original implementation is quite small though.

class KDTreeBoruvka(KDTreeBoruvkaBased):
    def desc(self) -> str:
        return "KD-Tree Boruvka"

    def setup(self, points: list[Point]):
        n = len(points)
        view = View(points)
        self._kd = build_kd_tree(view, Dimension(n=3))
        self._points = points

    def query_kd_tree(self, node, q, dim):
        if not node:
            return

        p = node.pivot
        Q, P = self.get_point(q), self.get_point(p)

        self.process_edge((q, p, dist(Q, P)))

        if node.is_leaf():
            return

        if Q[dim] <= P[dim]:
            main = node.left
            other = node.right
        else:
            main = node.right
            other = node.left

        self.query_kd_tree(main, q, dim.next())

        if other:
            lb = abs(Q[dim] - P[dim])
            ub = self.get_best_dist(q)
            if lb < ub:
                # needs to search in the other tree
                self.query_kd_tree(other, q, dim.next())

    def fill_best_edge_by_component(self):
        for i in range(len(self._points)):
            self.query_kd_tree(self._kd, i, Dimension(n=3))

Here, when comparing the query point q with a candidate (node.pivot), we call self.process_edge() which should ignore edges from points on the same component. We also work with indices of points instead of the points directly because we need the indices to do the union find properly. Other than that it’s exactly the same idea.

Optimizations

There are several optimizations we can do such as:

1-) Keep a bounding box for each node, so that instead of calculating the lower bound as the distance to the pivot line:

lb = abs(Q[dim] - P[dim])

We can do it to the bounding box of that subtree, which is tighter:

lb = point_box_distance(Q, other.bb)

2-) Store the set of the component ids in each subtree. If there’s only one and it’s the same as the query point, the entire search can be short-circuited.

This approach requires re-calculation after each iteration because the components change after merge but it’s a $O(N)$ step and the iteration is dominated by $O(N \log N)$ anyway. Then we can add this check in query_kd_tree():

class KDTreeBoruvka(KDTreeBoruvkaBased):
    def desc(self) -> str:
        return "KD-Tree Boruvka"

    def query_kd_tree(self, node, q, dim):
        if not node:
            return

        # prune: same component
        if node.components == {self.get_component(q)}:
            return

        ...

3-) A larger and more intrusive optimization is to not search the kd-tree for one point at a time, but instead traverse the kd-tree side by side and compute the distance for all points using a single traversal.

The algorithm is a lot more complicated and while my implementation of this version was 2x faster than the basic kd-tree one, it seems to be still a $O(N \log^2 N)$ process.

4-) Codex did a version with even more heuristics and optimizations, reaching a 4x speed up, but it removed a lot of the abstractions and the code became very hard to read, so I don’t think it’s very instructive.

The code for all these variants is on Github.

Experiments

During my experiments, I also used SciPy’s Minimum Spanning Tree algorithm, mostly to check correctness, but ended up including it in the benchmark as well. Since it’s so easy to translate code between languages, I asked Codex to convert its optimized Python implementation into C++ and Rust.

I ran these algorithms with sets of random points in 3D: $10^3$, $10^4$ and $10^5$ points. The runtimes in seconds are tabulated below:

We can see that C++/Rust are much faster than Python. The naive $O(N^2 \log N)$ C++ implementation is comparable to the optimized $O(N \log^2 N)$ ones even for $10^5$ and beats all the Python implementations for $10^4$.

Delaunay Triangulation

For the 2D case, a faster and simpler alternative exists, because it’s possible to show that the edges of a Euclidean MST of a set of points $P$ are a subset of the Delaunay triangulation of $P$.

The Delaunay triangulation of $N$ points in 2D has $O(N)$ edges and can be computed in $O(N \log N)$, so then running a regular $O(\abs{E} \log \abs{E})$ algorithm is very efficient. However for the 3D case, which is the case I was interested in, the number of edges can be $O(N^2)$.

So I studied Delaunay triangulation because I thought it was the best implementation for EMST but later I learned about the $O(N^2)$ edges.

Conclusion

This concludes a series of posts, which includes Delaunay triangulation, KD-trees and now the Borůvka algorithm.

It was motivated by a problem in Advent of Code which reduced to finding an EMST for 1,000 points in 3D. My Kruskal $O(N^2 \log N)$ ran in a few seconds but I wondered how much better it could be, so I went on this rabbit hole. Little did I know that if I had reimplemented Kruskal in C++, it would be faster than the Borůvka + KD-tree, but it’s never about the destination: I learned a lot during this process.

Related Posts

In An Introduction to Matroids we mentioned that the minimum/maximum spanning tree problem can be modeled as a matroid. Matroids can be solved by greedy polynomial-time algorithms and the Kruskal algorithm is the version that solves the matroid corresponding to MST. The Borůvka algorithm is also greedy.

There’s a generalization of the EMST called the Steiner Tree Problem, in which you can introduce intermediate nodes to try to reduce the cost of the solution, but solving this problem is NP-Complete.

The Steiner tree sounds a lot like the problem of Constructing Trees from a Distance Matrix which is also known as the tree metric realization problem, in which we can also introduce intermediate nodes, but in this case we’re deciding whether a tree can exist such that the distance between leaves (path length) matches a prescribed distance.

I’m using the term atomic post in reference to the concept of atomic habits. The idea is for posts to be as narrow as possible in scope, while still being self-contained.

Like with atomic habits, the idea is to lower the bar for posting in hopes of posting more often. I mentioned before the desire to write shower thoughts more often and that:

I still need a forcing function to make such posts, because otherwise I keep waiting until I have enough to write about them and it never happens.

In software development we strive for functions to be modular: short but also conceptually self-contained and with the right level of abstraction, so that they can be more easily reused. For posts, the same idea applies, but the ‘reuse’ is so that I can reference the idea in other posts and avoid repeating the same argument over and over.

A new motivation for doing this is that I started using LLMs to find connections between posts and it sometimes surprises me. My expectation is that the more content I have out the higher the chances of interesting connections.

I’ve set up some methodology to help with this. I keep Google docs for several topics I want to write about and whenever I read or think of something related, I add it to that document (this very post had a corresponding doc!). Whenever I find I have enough ideas collected, I will post them.

I’m also trying to be more deliberate on what non-ficcion books I read: I’ll pick topics I already plan to write about instead of random ones that happen to seem interesting to me at the time. The first example of this is Elbow Room by Daniel C. Dennett, since I’m interested in sharing my ideas on free will.

Once I have the notes, it should take me, ideally, less than an hour to finish the post. This way I don’t have to find a large block of time to write it. I won’t worry about being comprehensive, because I can always go back and edit the post.

Related Posts

In On Memory I described keeping book notes in Google Docs, but I’m gradually shifting to adding the notes under the a topic I’m interested in instead. Only when it’s something I hadn’t thought about or thoughts on the book/writing itself do I keep under the book note.

As early as 2021, in Writing Posts I mention I keep post drafts in Google Docs, but I do that when I already started the process of writing the post, while the new process is at an earlier stage, during ideation.

F14Map is Folly’s alternative to std::unordered_map and in this post we’ll explore this data structure in detail.

We’ll start with std::unordered_map which implements a more textbook version of hash maps using linked lists to handle collisions and then cover F14FastMap which uses a chunked (14 entries / chunk) open addressing implementation and leverages SIMD instructions for efficiency.

STL’s Hash Map

Recall that std::unordered_map is a hash map which uses linked lists for dealing with collision. Let’s cover the three main operations that can be performed: insertion, lookup and removal.

Insertion

When we insert a key-value pair, we first compute a hash h for the key to obtain a size_t (the hash function is configurable). Then we determine which bucket this key falls into by doing h % bucket_count. If two keys fall into the same bucket, we add to the bucket’s list.

In practice std::unordered_map uses a single linked list, but each bucket knows where its own list starts, for example:

sentinel ──> A ──> C ──> B ──> null
   ▲               ▲
   │               │
bucket[0]      bucket[2]     bucket[1] = nullptr

If we were to insert a new node for bucket 2, we can insert it between A and C and have bucket[2] point to this new node. The nodes store the value for the item.

Duplicates. We can’t just insert the item at the head and call it a day because if the key exists in the hash table we must update it instead of inserting, so we must traverse all the elements in the bucket to make sure if that’s the case. Since this is effectively what Lookup does, we’ll leave the details to that section.

In practice these operations get fused to avoid redoing work, for example determining the right bucket to insert is only done once.

Growing. As we insert more and more items in the hash table, the list for each bucket will grow and in the worst case degenerate to $O(n)$ insertion / lookup. So when the number of elements crosses a threshold, it increases the number of buckets by ~2x.

There are two “tracks” of thresholds, the default being primes that are roughly 2x apart:

2  5  11  23  47  97  197  397  797  1597  3203  6421  12853  25717

The other is power of 2. Primes are slower to apply % of (we cannot use bit shifting as with powers of 2), but they tend to hash more uniformly. In any case, every time the number of buckets changes, we need to rehash.

With a new bucket_count, the nodes will be reshuffled to different buckets so it reconstructs a new linked list from scratch by updating the next pointer of the nodes. This operation is done in $O(n)$. Note that the nodes themselves are not changed, so if we hold a reference to this node, it will be valid after the growing.

One way to avoid or reduce this overhead is to use .reserve(), much like we do for std::vector.

Lookup

For searching a key q, we need to traverse the whole list of a given bucket. How do we know when we “crossed” into the next bucket list? Each node stores the key hash (before the modular arithmetic), so when searching for an item it checks:

• Is the q hash equal to the node hash? If so, check for equality (==) on the key. If yes, return the value.
• If not, move forward. However, we also check if we crossed into the next bucket. We do so by doing % bucket_count: if it doesn’t match the one from q, it means it’s from a different bucket and we stop the search.

Note that we can’t store h % bucket_count in the node, because as we’ve seen, bucket_count can change.

Removal

Removal is straightforward, we search for an element via Lookup and remove the node from the linked list, potentially updating pointers on the bucket array.

The bucket_count never shrinks, even if most of the entries are removed.

F14Map

There are a few variants of F14Map and F14FastMap in particular decides which implementation to use at compile time based on the size of the item (key + value) being stored. If smaller than a threshold, it uses F14ValueMap which stores the items inline. Otherwise it uses F14VectorMap which externalizes the items to a vector. We’ll cover F14ValueMap first, then the differences with F14VectorMap.

F14Map relies on a lot of SIMD instructions which are not standardized between x86 and ARM architectures. For the sake of brevity we’ll use the x86 instructions, more specifically the 128-bit family SSE2.

As we did with std::unordered_map we’ll cover the 3 main operations – insertion, lookup and removal – at a high level, but then dive into some of the optimizations it uses that cannot be performed with the STL implementation, due to node allocation being done independently and hence fragmented in memory.

Insertion

The way F14Map handles collision is very different: first, the number of entries in a given bucket is fixed to 12 or 14 depending on the type of the item. The family of structures F14* is named after the latter value! I’m assuming 12 came later, probably due to some fine tuning for specific types. In any case, we’ll assume 14 for the remainder of the post for simplicity.

A structure holding the 14 entries is called a chunk. To determine which chunk an item should go to, we do h % chunk_count, much like in std::unordered_map. Note that chunk_count is always a power of 2. If that chunk is full, we try another chunk.

Before we explain how we search for the next chunk, we define the tag, which is an 8-bit value with the most significant bit set to 1 (we’ll explain the reason later) and the other 7 bits are extracted from the hash h, but the exact method for extraction varies.

We then calculate the stride, which is essentially 2 * tag + 1, which is always an odd number. This is the amount we’ll keep adding to the current index until we find a non-full chunk, wrapping around when it exceeds chunk_count. Because stride an odd number and chunk_count is a power of 2, they’re always co-prime, so we guarantee that by iterating by stride steps, we’ll cover all chunks before we repeat a chunk.

The fact that stride is dependent on the tag which has 7 bits of entropy should help make sure different keys have a different distribution of strides. In other words, the order in which chunks are visited for each item should be uniformly distributed, which should keep the expected number of chunks to visit until we find a valid chunk small.

First Free Position. It’s not enough to determine that a chunk is not full. We need to determine the position at which to insert the item in a chunk. This is the first optimization that leverages SIMD (Single-Instruction Multiple Data). The chunk structure has this shape:

struct F14Chunk<typename Item> {
    std::array<uint8_t, 14> tags_;
    uint8_t control_;
    uint8_t outboundOverflowCount_;
    std::array<Item, 14> rawItems_;
};

The variable control_ clubs 2 pieces of information, but for this post we only care about the highest 4 bits, hostedOverflowCount_, which counts “how many items currently living in this chunk don’t belong here?”. We increment it whenever we insert an item on a chunk that wasn’t its first choice.

The variable outboundOverflowCount_ counts “how many items wanted this chunk at any point in their search but it was full?”. We increment it whenever we try to insert an item in a chunk that was already full.

Note that these 2 variables pad the 14 bytes of tags_ into a 16-byte value which is the “unit” for the SIMD instructions it uses.

The data is stored in rawItems_. Each item in tags_ corresponds to an item in rawItems_ and it stores the tag we just discussed. If the position i is free, then tags_[i] = 0. Recall that tags have the most significant bit (MSB) set to 1. So if we want to find the first free position, we just need to find the first position that has MSB 0. This is done by this code:

uint16_t occupied = _mm_movemask_epi8(_mm_load_si128(&tags_[0])) & 0x3FFF;
uint16_t empty = occupied ^ 0x3FFF;
if (empty == 0) { /* chunk is full */ }
unsigned slot = __builtin_ctz(empty);

The function _mm_load_si128() loads 16 bytes into a 128-bit SIMD register. Since we’re passing the address of tags_[0] it will load the 14 bytes of tags_ and the 2-byte counters.

The function _mm_movemask_epi8() takes the MSB of each of the 16 bytes and packs that into a 16-bit integer. The mask & 0x3FFF clears the bits from control_ and outboundOverflowCount_ because we don’t care about them. Then ^ 0x3FFF flips the bits on the 14 LSB. So if all positions were occupied, we would have occupied = 0x3FFF and then empty = 0.

Finally __builtin_ctz() counts trailing zeros, so it returns the first bit of empty that is set, which was the bit in occupied that was 0 and thus the index of the first position that is free!

Growing. As in the std::unordered_map, once enough entries are inserted the structure must grow. Differently from the STL version though, the chunks and the entries they hold are not stable. They need to be moved to a different location of contiguous memory because the F14Chunk is stored as a contiguous array.

It always allocates a new array though because it can’t update the chunks in place. Note that this increases the memory temporarily, while both structures are live. Then it iterates over the existing chunks and items inside each to reinsert (std::move) them into the new chunks using the same strategy except that now chunk_count is twice the previous value.

This implies that the address of the items in the hash map changes during this process, so references are not stable, which is a semantic difference with std::unordered_map.

One optimization that can be used is that during this process tags_ is compact: there are no holes because there’s no removal until we finish the rehashing, so First Free Position can be found by keeping a counter of item count for the chunk.

Another optimization that can be done because we use tags_ as the check for the presence of an element, is that when we malloc a new array of F14Chunk we don’t need to memset() the array rawItems_, we can leave garbage there. We just need to clear the first 16 bytes:

std::memset(&tags_[0], '\0', 16);

which can be done with a single SIMD instruction.

Lookup

For lookup, first we compute the tag for the key and load into the SIMD register via _mm_set1_epi8(tag). This replicates the same byte 16 times into a register, for parallel comparison.

Then, for each chunk it visits, we load the tags into a register (like we do during insertion) then compare with the key’s tag using a single SIMD instruction and then find the position where the bit is set. The code is roughly:

// once per search
__m128i needleV = _mm_set1_epi8(tag);

// for each chunk
auto tagV = _mm_load_si128(&tags_[0]);
auto eqV  = _mm_cmpeq_epi8(tagV, needleV);
auto mask = _mm_movemask_epi8(eqV) & 0x3FFF;
if (!mask) {
    // tag not in chunk
}
unsigned slot = __builtin_ctz(mask);

The function _mm_cmpeq_epi8 returns a 16-byte register. A byte is set to 0xFF if the corresponding bytes of the input are equal or 0x00 otherwise. Recall that _mm_movemask_epi8() takes the MSB of each of the 16 bytes and packs that into a 16-bit integer. Again we need 0x3FFF to exclude the bits from control_ and outboundOverflowCount_.

This acts as an early filter (think of a simplified bloom filter), but we still need to check if the key corresponding to that tag matches the searched one, so we’d check rawItems_[slot]. Note that mask might have multiple bits set if different keys map to the same tag and are on the same chunk, so we need to iterate.

The outboundOverflowCount_ can be used to stop searching early: if it’s 0, then it means no entry tried to insert on this chunk and could not, including the item we’re looking for, so if we don’t find the key here, there’s no point in continuing the search. Without this check we’d always have to scan all $O(n)$ chunks to be sure a key doesn’t exist anywhere.

Removal

The first part of the removal consists of finding the location of the element and this is exactly like Lookup. The second part consists of the cleanup. First we clear the corresponding byte in tags_ and destroy the value in rawItems_.

Then we need to update outboundOverflowCount_. We do this by replaying the insertion of the key being removed. It will start at the initial chunk (which depends only on the hash of the key) and then we’ll move by strides until we reach the chunk we’re at. On every chunk we visit we subtract from outboundOverflowCount_. The chunks we visit during removal are the same we did during insertion, no matter if other items were inserted or removed in between!

Removal by Iterator. The flow described above is when we have the key and want to remove it. This involves a lookup to find the position before erasing.

Another case is when we already have the iterator, in which case we can go straight to the removal. This is useful to apply a “filter” over the hash table, for example:

for (auto it = m.begin(); it != m.end(); ) {
    if (pred(*it)) it = m.erase(it);
    else ++it;
}

Here .erase(it) returns an iterator to the next item. When we have the iterator we know the chunk and index to remove but we don’t have the hash nor the tag. To perform the loop to update outboundOverflowCount_ we’d need to compute the hash from the key we’re at.

This is where the variable hostedOverflowCount_ can help us if it’s 0. Recall it counts: “how many items currently living in this chunk don’t belong here?”. If it’s 0, it means this is the first chunk and we don’t need to look further and we’d not need to compute the hash!

Otherwise, we need to decrement hostedOverflowCount_ after removing the element.

F14VectorMap

We use F14ValueMap whenever sizeof(std::pair<Key, Value>) < 24. Otherwise we do F14VectorMap. The main difference is that the item is not stored in the chunk’s rawItems_, but instead stored in an array, roughly:

std::pair<const Key, Value>* values_{nullptr};

which is malloc()‘ed to have chunk_count entries, then rawItems_ only stores indices/offsets to values_.

Insertion. This is very similar to before, except that we now store the value in values_[size], store size in rawItems_ and then increment size. When we need to grow, we malloc() a new array of size chunk_count * 2, and move the old values there.

Removal. The values_ array is dense, meaning that the first size positions are always occupied and the remaining contain garbage. This invariant can be maintained when removing position i by moving the last element to there and decrementing size by one. Then the rawItems_ corresponding to the swapped item must be updated to point to i. The way it’s done is by taking the key from that item to find the chunk + index using the same process as the insertion/lookup.

Comparison

The major difference between std::unordered_map and F14FastMap is that the former allocates a node for each entry, whereas for F14 it stores them in contiguous chunks of memory (both for the F14Chunk and values_ array). The STL version performs one malloc() per entry while the F14 only does malloc() when growing.

The tag system from F14 allows leveraging SIMD instructions and because the chunk is aligned at 16 bytes:

template <typename ItemType>
struct alignas(16) F14Chunk {
    std::array<uint8_t, 14> tags_;
    ...
};

The tags_ fit in the cache line.

Conclusion

In this post we learned how std::unordered_map and folly::F14FastMap work behind the scenes! I’m sure I got a bunch of details wrong but I have a cohesive understanding of the overall implementation. I’ve been curious about the 14 since I started working with folly::F14FastMap and it now all makes sense!

It was a delight to study folly::F14FastMap: both the data structure and the low-level optimizations are very clever!

Related Posts

We mentioned the post CPU Cache for cache lines, but also that post describes the implementation of a cache line as a hash table!

In Velox: The Vector we discussed the dictionary-encoding in which we only store indices to the actual underlying data, similar to F14VectorMap. In Velox: Memory we mentioned that memory is never given back to the OS, which is the same behavior for F14VectorMap, it never shrinks the number of buckets. It makes sense statistically to not shrink given that it’s expensive and that, without more information, we can assume we’ll need the peak memory in the future.

The post HyperLogLog in Rust also uses a hash function to determine the bucket and a function first_non_zero_bit_position which is equivalent to __builtin_ctz().

References

• [1] Claude Code
• [2] Github - facebook/folly

Velox is an open source C++ library by Meta that can be used to perform computation common to distributed engines like Presto.

Its offerings include columnar operations, a rich type system, an expression parser and a smart resource management such as memory [1]. In this series of posts we’ll go over different components of Velox.

In this post we’ll study memory management done by Velox.

Previous posts on the series:

• Vectors
• A Simple Application
• UDFs

Velox has a custom memory allocator system because the same process can run short-lived queries from different clients. To avoid one query causing others to OOM, it cannot rely on OS-level limits such as cgroups and thus uses its own arbitrage.

Components

There are 4 main components involved in memory allocation: the manager, the allocator, the arbitrator and the pools. The pools themselves are subdivided into root pool, “inner” pool and leaf pool.

The memory allocator (MemoryAllocator) is the one interfacing with the OS and making the actual allocation.

The memory arbitrator (MemoryArbitrator) is an interface, and is implemented by classes like SharedArbitrator. Its purpose is to distribute quotas to the pools, but it doesn’t perform the memory allocation.

Memory pools (MemoryPool) are organized in a tree hierarchy: at the top there is the root pool, in between there are the inner pools and at the leaves the leaf pools. The root and inner pools are of aggregate type, meaning they are logical grouping of the leaf pools and they’re mostly used to keep aggregated statistics such as total memory used by the subtree. The leaf pool is the one Velox internals interface with.

The idea is that each query corresponds to a root pool. A given process can be executing a number of queries. Each query runs one or more tasks, each having a pool as a child of the query root pool. Then for each node in the plan (see Velox Application) we have another pool. Each node might be converted into one or more operators and one instance of each operator will exist for each driver (in the parallel mode). Finally, each such instance of the operator has the leaf pool.

The memory manager (MemoryManager) is the singleton which can be used to construct the allocator and pools.

Allocation

Allocation typically happens at the leaf pool and it might bubble up all the way to the memory arbitrator. Let’s go through the flow.

Each leaf pool tracks how much memory it’s using (we’ll name it used) and how much it has allocated to it (we’ll name it reserved), and both start at 0. When the caller requests requested amount of memory from the leaf pool, it will check if used + requested <= reserved. If yes, then it just updates used += requested.

If used + requested > reserved, it will ask the root pool to increase its reservation. The request goes through its parents because then they can update the aggregated reserved values. Once it reaches the root, it will check if it’s within its limit.

The root pool has two relevant properties: capacity and maxCapacity. The capacity is analogous to the leaf pool’s reserved: if the root pool exceeds this limit, it requests the arbitrator to increase its capacity. The maxCapacity is a limit on capacity. It’s a static amount that is set when creating the root pool (unbounded by default). If the increase in capacity would exceed this value, the pool will try reclaiming memory (see Reclamation).

The arbitrator also contains a limit, capacity, which is set by the application. If a request for more memory would exceed this value, it will attempt memory reclamation (see Reclamation). Note that the arbitrator limit is global.

So far, all this escalation from leaf pool → root pool → arbitrator is accounting. It checks if the allocation can be performed, but the actual allocation happens via the memory allocator and is initiated by the leaf pool.

The memory allocator also has a limit, also set by the application. If this limit is exceeded, it will result in the error VELOX_MEM_ALLOC_ERROR. Note that memory allocation can be requested by callers directly to the allocator, without using pools.

The arbitrator capacity must be less than that of the allocator. It’s not exactly the same because of the last point: not all memory allocated is accounted for by the arbitrator, only the ones requested from the pools.

Custom Allocator vs. Malloc

The allocator has 2 flavors: malloc or mmap, which uses malloc() and mmap() to allocate memory, respectively. Malloc is the default mode.

Recall that malloc() can itself be customized by linking a library such as jemalloc during compilation. As we’ve seen libraries like jemalloc do a lot of the heavy lifting on memory allocation, balancing performance and efficiency (reduce fragmentation) and making use of modern hardware and OS (e.g. memory layout and multi-thread). So for this flavor, the Velox memory allocator mostly delegates to malloc().

For the mmap() version, since it requests larger pages to the OS, Velox has to add its custom logic to do the fine grained allocation and we won’t cover it here.

One of the major downsides of this allocator compared to jemalloc is when there are a lot of threads, contention can be significant. As we’ve seen, jemalloc avoids this by using arenas (see Multi-threads and Arenas in jemalloc).

Deallocation

Freeing memory is almost the exact opposite of allocation. When called on a leaf pool, it reduces used. If the amount of free memory (reserved - used) crosses a threshold (the quantized sizes), then reserved is reduced and this amount flows back to the root pool.

Reclamation

As we discussed, there are 2 cases in which memory reclamation can happen:

First is when the maxCapacity of a root pool would be exceeded, in which case it will choose a leaf pool from which it can reclaim memory. First it checks if the pool can reduce its reserved and give it back to the root. If not, it tries to find a leaf pool in its subtree to spill its state to disk. Once an operator spills to disk its used memory drops and it follows the Deallocation process. Only certain operators can do spilling (e.g. HashBuild, OrderBy, etc.).

The second case is when the capacity on the arbitrator is reached. It will choose a root pool from whom to reclaim memory and then the process is the same as above.

If reclamation cannot be done, a query gets killed. In the first case it’s the query making the request; in the second the arbitrator will select the query to be killed. After reclamation is done, the arbitrator can re-allocate the recouped memory to fulfill the original request.

Other

Cache

Velox supports cache by building on top of the memory framework. The component is called AsyncDataCache and it allocates from the MemoryAllocator directly (not through pools).

This cache has a secondary (optional) layer backed by SSD. When enough bytes have been stored in memory, it eventually flushes to SSD. Data is kept in memory as well, but when memory reclamation arrives those can be easily dropped.

Eviction happens when no more data can be stored in RAM, which is the same trigger we’ve seen for the memory allocation.

Conclusion

In this post we’ve learned about the memory management done by Velox. It was one of the most mysterious parts of it, and I think I have a much better grasp on it. It’s not as complex as I imagined, but it’s a lot more complicated than I though with lots of components.

Writing also provided some insights on the motivation, for example, being designed to run multiple independent queries in the same process, which seems to be modeled after the Presto query engine (which is the system it aimed to replace).

It also clarified my understanding of memory pools and the two underlying allocator systems, malloc and custom.

Related Posts

We already discussed jemalloc which itself is a user-space, custom memory allocator that internally implements Buddy Memory Allocation.

The summary [Book] Systems Performance mentions that allocators don’t return memory to the kernel once free is called, but keep them around. Another relevant behavior is that the virtual to physical mapping for memory is done lazily by the OS: it only happens when memory is written to. In a sense allocation and reclamation in Velox are also lazy.

In my last post about KD-tree, I vibe-coded a JavaScript application to demonstrate visually, step-by-step, how to construct and query a KD-tree. Codex one-shotted it in a few minutes; I spent a few minutes linking it to the blog post but did not even take a look at the code.

I’ve been using AI more and more both at work and for personal stuff but I also deliberately avoid using AI for specific tasks. I wanted to spend some time reflecting on this and document it.

Writing

For Work

Since the beginning of the year, I write effectively 0% of my code by hand at work. I don’t consider it vibe-coding because I spend a large amount of time reading the code and making sure I understand what it does. I do skim over details and most of the generated unit tests.

AI sometimes writes lengthy comments with too much context which make it cumbersome to understand. If it’s too much I manually erase chunks and sometimes rewrite it myself. I’m very close to the point of also not reading existing code directly so in theory I should not worry about walls of text which could be useful for AI.

I also write PR summaries manually, mostly to force myself to understand what the PR does. I found many times I had a gap in understanding of a change while drafting the PR summary.

For Personal

For learning algorithms or solving programming puzzles, I write 100% of the code. I even turn off AI auto-complete. One reason is to keep my programming skills sharp (it’s like hiking instead of driving), even if I might not need them in the future. The more important reason is that I find I learn better when I do things myself.

That’s also the reason why I write 100% of my blog posts manually. I use AI to correct typos and grammatical errors but explicitly ask it not to rephrase things, even if it improves clarity. I want my blog to reflect my way of thinking, even if it’s gradually but increasingly being influenced by AI.

A few times I did ask for feedback from AI on how to convey my ideas better, especially when reviewing books, which helped clarify my understanding.

Verum Ipsum Factum

This is a quote by Giambattista Vico, an Italian philosopher, and it translates to We only know what we make. I read this in the book Building a Second Brain and this is the main motivation to keep doing things manually and that’s the main reason I blog.

In college, I recall we had a physics class that consisted in running experiments and using machines such as an oscilloscope. We worked in groups of 3, and each experiment we rotate on who did the manual work and who did the note taking, etc.

There was this one experiment with complicated steps and I wasn’t the one doing the manual work. I tried as hard as I could to memorize how my classmate did it, but wasn’t sure I had internalized it. As chance would have it, this was exactly the experiment we were asked to reproduce in the final exam and I couldn’t do it properly. Luckily I didn’t fail the course but it’s a lesson I’ll always remember.

Reading

I’ve been relying more and more on AI for learning, instead of reading textbooks or Wikipedia. I still love reading books for fun, but for difficult subjects I find it more effective to do everything through AI. I have specific learning patterns and some textbooks are so dense I waste time trying to decipher missing information (I spent 2 years reading a book on Complex Analysis).

AI is very useful for digging into details on demand too. I might not care about details of a specific topic but want to keep asking for details on others. Printed books often cannot provide such a tailored experience. AI makes it easier to follow my curiosity. If I had to buy textbooks or spend hours researching, I’d just not do it.

The one concern with using AI is that some say that true learning only comes with struggle. If AI makes it too easy to digest things, am I learning optimally? Writing posts explaining stuff (even if no one reads it!) is even more important now, because it tests my understanding and keeps me honest, but I don’t know if it’s enough.

Principles

My guiding principle for when to use AI vs. not: what’s the most effective way to learn and understand? I don’t have qualms about relying on AI to do 100% of my grunt work, but I also don’t aim to blindly automate my work without learning anything in the process.

I’m fine to blindly let AI build tools to advance my projects but I’d like to be on top of the architecture and the algorithms that it uses, trade-offs taken, how the system deals with corner cases, etc.

I still think it’s valuable to know and understand things instead of being an AI puppet, but I also recognize that the value of knowledge might shrink the more AI can expand the scope of what it can do. I’d still want to learn for fun, even if it doesn’t matter, the same way I enjoy reading biology books even though I likely won’t make use of it. I have more thoughts on this but I’ll leave it for its own post.

Conclusion

To circle back to the start of the post, learning as a goal makes it easy to justify vibe-coding the JavaScript app to better understand Delaunay triangulation and KD-trees. My goal is not to learn how to write JavaScript apps, it’s to learn the algorithms themselves.

Related Posts

In my 2025 in Review post, I mentioned how I found learning history from Wikipedia instead of books more effective. I’ve since switched to using AI entirely for it. Most of my notes for History have been through sessions with ChatGPT.

In The Weierstrass ℘-Function also mention:

This proof is also not provided in Ahlfors (in any obvious way at least) and I relied entirely on ChatGPT to understand it! Like with learning history, I’m finding using ChatGPT a lot more effective at learning math: I can ask for it to explain me things in different angles and dig into different parts.

This has been a consistent theme in many of my recent posts, so now that I have a dedicated post on this point I can stop repeating them or just link to here.

On Documentation I asked the question: Who Should Write (documentation)?. Now we have a new option: AI. We should also consider a new point: who is the audience of documentation now? Humans or AI? Worth a separate post when my thoughts on this mature.

Jon Bentley is an American computer scientist, famously known for his book Programming Pearls. He also came up with the data structure called k-d tree while an undergrad at Stanford and published it in a paper titled Multidimensional Binary Search Trees Used for Associative Searching in 1975.

In this post we study the kd-tree data structure, how to construct and perform queries on it. Then we provide an implementation in Python.

Nearest Neighbors

The “kd” in kd-tree stands for k-dimensional tree. We can think of it as a generalization of a binary search. In one version of the binary search problem we are given an array and a query value $q$. We need to find the element in the array that is the closest to $q$. This is called nearest neighbors search.

The problem kd-tree solves is essentially the same but the values are points in a k-dimensional space instead of a scalar (1d). The variant we’re interested in is the static one, i.e. we don’t allow inserting, updating or removing points.

To solve the 1d case, we consider the points in the line and let $m$ be the median point. If $m < q$ then we can discard half of the points from the input, those having values less than $m$. If $m > q$, we do the same for the other half. If $m = q$, we found the closest point. This property allows us to find the nearest neighbor in 1d in $O(\log n)$ time by keeping the points sorted.

For the kd case, there’s no single ordering, so what can we do?

Partition

One way to see the search space of the 1d case is that of partitions organized as a binary search tree (BST). First we partition the points into 2: those with value less or equal than the median $m$, and those with value greater than $m$. Each partition has roughly the same number of points. We then subdivide each partition the same way, recursively, so we end up with a BST.

For the kd-tree we use the same idea but we “alternate” on which dimension to partition on. This is most intuitive for 2d: first we partition on the $x$ axis, i.e. we split the points by finding the median of the $x$ coordinate $x_m$ and have one partition contain all points with $x \le x_m$ and another the points with $x \gt x_m$.

For the next level we partition on the $y$ axis instead. Geometrically each node in the 2nd level of the BST represents a quadrant containing roughly 1/4 of the points.

Searching

Unfortunately, for dimensions greater than 1, knowing which partition a query point $q$ belongs to does not allow us to discard the other half. Here’s a counter-example:

However, once we do find the partition of $q$, we do get an upper bound for the distance of the closest neighbor. This helps prune searches to the other partitions if there is not a lot of degeneracy.

So the idea of the algorithm is: traverse the BST to find the right partition and note the distance $d$ from $q$ to the point on that partition. As we backtrack, we ask “could this alternative branch contain the closest point?”.

We can quickly determine the lower bound between $q$ and any point of that branch by computing the distance between $q$ and the hyperplane (or line in the 2d case) which is $\abs{k_q - k_m}$ where $k$ is the dimension chosen for the partition and $k_m$ the median point. If this value is higher than $d$ we don’t need to keep searching. Otherwise we repeat the process there.

Once we reach the leaf on that subtree, we’ll potentially update $d$ and repeat the exact same process. In the worst case, we might end up visiting each leaf in this tree which would lead to a $O(n)$ complexity.

Implementation

We’ll implement a static, general dimension, kd-tree in Python. First we define some helper classes.

Data Structures

We start with Point, which is a thin wrapper on top of tuple, that implements some utility operators:

@dataclass(frozen=True)
class Point:
    data: tuple[int, ...]

    def __getitem__(self, dim: Dimension):
        return self.data[dim.v]

    def __len__(self) -> int:
        return len(self.data)

    def __abs__(self) -> float:
        return sqrt(sum(x*x for x in self.data))

    def __sub__(self, other: Point):
        sub = tuple(a-b for a,b in zip(self.data, other.data))
        return Point(sub)

Where Dimension is just a simple wrapper to carry around the max dimension so we can get the next dimension in modular fashion:

@dataclass(frozen=True)
class Dimension:
    n: int
    v: int = 0

    def next(self):
        next_v = (self.v + 1) % self.n
        return Dimension(self.n, next_v)

Since the kd-tree is a binary tree, the node is straightforward. We also add an is_leaf() to help with readability:

@dataclass
class KDNode:
    pivot: Point
    left: KDNode
    right: KDNode

    def is_leaf(self):
        return self.left is None and self.right is None

Finally we define a View class representing a sub-range, or a view, of a list. This is because we plan to avoid copying the array and instead work with sub-ranges of it:

@dataclass
class View:
    pts: list[Point]
    lo: int = 0
    hi: int | None = None

    def __post_init__(self):
        if self.hi is None:
            object.__setattr__(self, 'hi', len(self.pts) - 1)

    def range(self, lo, hi):
        return View(self.pts, lo, hi)

    def swap(self, i, j):
        self.pts[i], self.pts[j] = self.pts[j], self.pts[i]

    def __getitem__(self, i: int):
        return self.pts[i]

    def __setitem__(self, i: int, p: Point):
        self.pts[i] = p

Median

To construct the kd-tree we need to find the median point, i.e. a point such that half of the points in the set are smaller than it, and half are greater for a particular dimension. We can use the quick-select algorithm for this:

• Choose a random element from the array. Call it pivot
• Rearrange the array so that all elements smaller than pivot appear before it and all elements greater than it appear after.
• After this, if pivot is towards the beginning of the array, repeat for the right side of the array. If towards the end, repeat for the left side. The goal is to eventually have pivot appear at the middle of the array.

The partition() function below implements the second bullet point:

def partition(pts: View, p: int, dim: Dimension) -> int:
    """
    Lomuto partition. Rearrange pts[lo:hi] so that
    if i < p, then pts[i][dim] <= pts[p][dim]
    if i > p, then pts[i][dim] > pts[p][dim]
    """
    lo, hi = pts.lo, pts.hi
    pivot = pts[p]
    pts.swap(p, hi)
    store = lo

    for i in range(lo, hi):
        if pts[i][dim] <= pivot[dim]:
            pts.swap(store, i)
            store += 1

    pts.swap(store, hi)
    return store

The first and third points are implemented by the outer function:

def median_index(pts: View, dim: Dimension) -> int:
    lo, hi = pts.lo, pts.hi
    k = lo + (hi - lo) // 2
    while True:
        if lo == hi:
            return lo
        p = partition(pts.range(lo, hi), random.randint(lo, hi), dim)
        if p == k:
            return p
        elif p > k:
            hi = p - 1
        else:
            lo = p + 1

It’s possible to show this algorithm is $O(n)$ on average but $O(n^2)$ in the worst case.

Building

Open the interactive applet that builds a kd-tree.

Once we have the function to find the median, building the kd-tree is very simple:

def build_kd_tree(pts, dim):
    if pts.lo > pts.hi:
        return None

    mi = median_index(pts, dim)

    left = build_kd_tree(pts.range(pts.lo, mi - 1), dim.next())
    right = build_kd_tree(pts.range(mi + 1, pts.hi), dim.next())

    return KDNode(pivot=pts[mi], left=left, right=right)

Searching

Open the interactive nearest-neighbor search.

Search is relatively simple too, when traversing a node:

• Check if the current pivot is the closest point.
• Recurse on the query point’s partition.
• If the lower bound from the other partition is smaller than what we have so far, recurse on the other partition as well.

def query_kd_tree(node, q, dim, ub = inf):
    if not node:
        return None

    p = node.pivot

    if node.is_leaf():
        return p

    if q[dim] <= p[dim]:
        main = node.left
        other = node.right
    else:
        main = node.right
        other = node.left

    c = None

    # pivot is a candidate
    if dist(q, p) < ub:
        c = p
        ub = dist(q, p)

    # main branch
    c1 = query_kd_tree(main, q, dim.next(), ub)
    if dist(q, c1) < ub:
        c = c1
        ub = dist(q, c1)

    # minimum distance we can expect to find
    # on the other branch
    lb = abs(q[dim] - p[dim])
    if lb < ub:
        c2 = query_kd_tree(other, q, dim.next(), ub)
        if dist(q, c2) < ub:
            c = c2

    return c

In the best case the search behaves like binary search and it takes only $O(\log n)$ operations but in degenerate cases it can be $O(n)$ (i.e. all nodes of the tree need to be searched).

The full code is on Github.

Experiments

I ran this code with random inputs of 10,000 points with 100 query points. The average number of nodes visited was 34.59 ± 15.61 which is logarithmic on the input size.

Conclusion

Kd-tree is on my list of data structures or algorithms I had heard of, possibly used some existing implementation, but never took the time to study in detail. Another recent example is the Delaunay Triangulation.

One thing I skipped for this post is any theoretical analysis that shows how kd-trees perform on average.

The JavaScript applet was vibe-coded using Codex, which one-shot it (after finding a bug in my implementation) in less than 5 minutes, with this prompt verbatim:

I want to build a JavaScript self-contained app that shows how the kd-tree algoritm works. I’m going to use this implementation: /Users/kunigami/workspace/python/kd-tree/kd-tree.py . The idea: use svg to draw a 512 x 512 box, generate 13 random points. It should render the dividing line from the pivot. On the left, render a binary tree. Inside the node show the pivot point. Ok? We want to do this interactive, so we should have a parameter step. the rendering should go only up to step N while doing a pre-order ‘traversal’. Then we add buttons prev / next which increment N and re-render. Feasible?

This was provided more or less on the spot: you can see I added details after such as the buttons, so it wasn’t a well thought out prompt. I was astonished by the result. I have been using Claude Code to generate code and sometimes it does generate entire code features for me, but maybe because this one was visual, I felt more impressed!

I’m pretty sure there are many such applets out there from which the model can learn, so it probably helped in it being so successful.

It would have taken me hours to implement this and I wouldn’t have found it worth spending the extra time. With AI I can not only enrich my posts with these visualizations but build these while studying to better understand things.

It also built the search version of it in one shot. I need to raise my expectations for what LLMs can accomplish. In Delaunay Triangulation I said:

It did cross my mind to write a JavaScript-based on for the Bowyer-Watson algorithm, especially if Claude can do most of the implementation, but it would still take some work to test, polish and read the code, so I decided to pass.

Related Posts

Finding the median of a set of points efficiently is an interesting problem. T-Digest in Python tackles this statistical problem from a different angle: how to compute the P50 (and other percentiles) online, i.e. without storing all points explicitly.

Binary search trees, such as the red-black tree, implement the 1d version of the kd-tree but allow for efficient insertion and removal. In Consistent Hashing we use Rust’s implementation of a red-black tree to efficiently search the closest server of a given key in the “ring”.

Skip Lists is a probabilistic alternative to binary search trees for range search. A natural question is: can skip lists be generalized to k dimensions? Eppstein et al. have papers on those: The Skip Quadtree and Skip-webs.

Another multi-dimensional probabilistic structure for nearest neighbor search is called Hierarchical Navigable Small World, which is apparently very popular for vector search, an important aspect of modern ML.