Problem 98 · Data Engineering & Streaming

LSM Tree: Memtable Flush and SSTable Compaction

HardClassic DSA & systems

01Problem

Build a tiny LSM store: writes go to a memtable that flushes to immutable sorted SSTables when full; reads check newest first; deletes write tombstones; compaction merges SSTables keeping the newest version and dropping tombstones.

02Key insight

LSM trades read amplification for sequential writes. Correctness hinges on recency ordering: memtable > newer SSTables > older ones, and tombstones must shadow older values until a full compaction reaches the bottom level, where they can be dropped. Compaction is a k-way merge keyed by (key, -age).

03Solution

solution.pypython52 lines

import heapq
from bisect import bisect_left
from typing import Any, Optional

TOMBSTONE = object()
_MISSING = object()  # private sentinel: a stored None is a real value
SSTable = tuple[list[Any], list[Any]]  # (sorted keys, values)


class LSMTree:
    def __init__(self, memtable_limit: int = 4) -> None:
        self.memtable_limit = memtable_limit
        self.memtable: dict[Any, Any] = {}
        self.sstables: list[SSTable] = []  # oldest first, newest last

    def put(self, key: Any, value: Any) -> None:
        self.memtable[key] = value
        if len(self.memtable) >= self.memtable_limit:
            self.flush()

    def delete(self, key: Any) -> None:
        self.put(key, TOMBSTONE)

    def get(self, key: Any) -> Optional[Any]:
        value = self.memtable.get(key, _MISSING)
        if value is _MISSING:
            for keys, values in reversed(self.sstables):  # newest wins
                i = bisect_left(keys, key)
                if i < len(keys) and keys[i] == key:
                    value = values[i]
                    break
        return None if value is TOMBSTONE or value is _MISSING else value

    def flush(self) -> None:
        if self.memtable:
            items = sorted(self.memtable.items(), key=lambda kv: kv[0])
            self.sstables.append(([k for k, _ in items], [v for _, v in items]))
            self.memtable = {}

    def compact(self) -> None:
        """Full compaction: merge all SSTables, keep newest version, drop tombstones."""
        runs = [[(k, -age, v) for k, v in zip(keys, vals)]
                for age, (keys, vals) in enumerate(self.sstables)]
        keys_out, vals_out, last = [], [], object()
        for key, _, value in heapq.merge(*runs, key=lambda t: (t[0], t[1])):
            if key == last:
                continue  # older version of a key already emitted
            last = key
            if value is not TOMBSTONE:
                keys_out.append(key)
                vals_out.append(value)
        self.sstables = [(keys_out, vals_out)] if keys_out else []

Try it first.

Timeput O(1) amortized (+O(M log M) flush); get O(T log N); compact O(N log T)

SpaceO(N)

tests.py · 4 passing

tests.pypython12 lines

db = LSMTree(memtable_limit=2)
db.put("a", 1); db.put("b", 2)       # flush -> sst0
db.put("a", 10); db.delete("b")      # flush -> sst1 (newer)
db.put("c", 3)
assert (db.get("a"), db.get("b"), db.get("c"), len(db.sstables)) == (10, None, 3, 2)
db.flush(); db.compact()
assert len(db.sstables) == 1 and db.sstables[0] == (["a", "c"], [10, 3])
assert db.get("b") is None and db.get("a") == 10
db2 = LSMTree(memtable_limit=2)
db2.put("a", 1); db2.put("b", 2)     # flush: a=1 on disk
db2.put("a", None)                   # stored None in memtable shadows a=1
assert db2.get("a") is None and db2.get("zz") is None

04Staff-level follow-up

How do you make reads fast and writes durable? Per-SSTable Bloom filters skip tables that can't contain the key, a WAL makes the memtable crash-safe, and leveled (vs size-tiered) compaction trades write amplification for bounded read/space amplification.