# /// script
# requires-python = ">=3.10"
# dependencies = ["rocksdict", "lmdb"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""LSM against B-tree: write amplification.

Bytes written to the device / bytes of user data, for RocksDB (an LSM) and
LMDB (a B-tree), with small rows and with 8 KB rows; then RocksDB with
key-value separation (blob files) for the large rows. Device bytes come from
/proc/self/io (Linux), so this lab runs in a Linux container.
"""
import argparse, json, os, random, shutil, tempfile

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--small-rows", type=int, default=5_000_000, help="Small rows (the chapter uses 50 million).")
p.add_argument("--small-bytes", type=int, default=100, help="Value size of a small row.")
p.add_argument("--large-rows", type=int, default=200_000, help="Large rows (the chapter uses 1 million).")
p.add_argument("--large-bytes", type=int, default=8192, help="Value size of a large row.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
random.seed(args.seed)

def written():
    for line in open("/proc/self/io"):
        if line.startswith("write_bytes"): return int(line.split()[1])

def rocks(path, n, size, blob=False):
    from rocksdict import Rdict, Options, WriteOptions
    o = Options(); o.create_if_missing(True)
    if blob: o.set_enable_blob_files(True); o.set_min_blob_size(1024)
    db = Rdict(path, o)
    for i in range(n):
        db[random.randbytes(16)] = os.urandom(size)                       # random keys: real compaction work
    db.flush(); db.compact_range(None, None); db.close()

def lmdb_store(path, n, size):
    import lmdb
    env = lmdb.open(path, map_size=64 << 30, sync=True)
    for start in range(0, n, 1000):                                        # one transaction per 1,000 rows
        with env.begin(write=True) as tx:
            for i in range(start, min(start + 1000, n)): tx.put(random.randbytes(16), os.urandom(size))
    env.sync(); env.close()

rows = []
for label, fn, n, size, kw in (("RocksDB (LSM), small rows", rocks, args.small_rows, args.small_bytes, {}),
                               ("LMDB (B-tree), small rows", lmdb_store, args.small_rows, args.small_bytes, {}),
                               ("RocksDB (LSM), 8 KB rows", rocks, args.large_rows, args.large_bytes, {}),
                               ("LMDB (B-tree), 8 KB rows", lmdb_store, args.large_rows, args.large_bytes, {}),
                               ("RocksDB + blob files, 8 KB rows", rocks, args.large_rows, args.large_bytes, {"blob": True})):
    d = tempfile.mkdtemp(dir="/tmp")
    os.sync(); before = written()
    fn(d, n, size, **kw)
    os.sync(); amp = (written() - before) / (n * (size + 16))
    rows.append({"store": label, "write_amplification": amp}); print(f"{label:<34} write amplification {amp:6.1f}x")
    shutil.rmtree(d)

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"stores": rows}}, f)