# /// script
# requires-python = ">=3.10"
# dependencies = ["duckdb", "numpy", "pyarrow", "pyroaring"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Copy-on-write against delete vectors.

A local Parquet table of 640,000-row files. 5,000 random row corrections are
applied two ways: (a) copy-on-write, rewriting every file a correction
touches; (b) merge-on-read, writing one small delta file (deleted positions
plus the corrected rows) per batch of corrections. Bytes written by each, and
the time of a full scan as unmerged delta files accumulate (1, 10, 100). Then
a Roaring bitmap of 1,000 deleted positions in one file.
"""
import argparse, json, os, shutil, tempfile, time
import duckdb, numpy as np, pyarrow as pa, pyarrow.parquet as pq
from pyroaring import BitMap

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--files", type=int, default=8, help="Data files.")
p.add_argument("--rows-per-file", type=int, default=640_000, help="Rows per file.")
p.add_argument("--corrections", type=int, default=5000, help="Row corrections.")
p.add_argument("--deltas", default="1,10,100", help="Unmerged delta files to scan through.")
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()

rng = np.random.default_rng(args.seed); d = tempfile.mkdtemp(); N = args.files * args.rows_per_file
for f in range(args.files):
    ids = np.arange(f * args.rows_per_file, (f + 1) * args.rows_per_file)
    pq.write_table(pa.table({"id": ids, "amount": rng.integers(1, 10_000, len(ids)), "note": np.char.add("n", (ids % 9973).astype(str))}),
                   os.path.join(d, f"base-{f}.parquet"))
base_bytes = sum(os.path.getsize(os.path.join(d, x)) for x in os.listdir(d))
fix = np.sort(rng.choice(N, args.corrections, replace=False))
# (a) copy-on-write: every touched file is rewritten
touched = sorted(set((fix // args.rows_per_file).tolist())); t = time.perf_counter(); cow_bytes = 0
for f in touched:
    tb = pq.read_table(os.path.join(d, f"base-{f}.parquet")); amt = tb["amount"].to_numpy().copy()
    local = fix[(fix // args.rows_per_file) == f] - f * args.rows_per_file; amt[local] = 0
    out = os.path.join(d, f"cow-{f}.parquet"); pq.write_table(tb.set_column(1, "amount", pa.array(amt)), out); cow_bytes += os.path.getsize(out); os.remove(out)
cow_s = time.perf_counter() - t
# (b) merge-on-read: one small delta file per batch
def write_deltas(k):
    for old in [x for x in os.listdir(d) if x.startswith("delta-")]: os.remove(os.path.join(d, old))
    total = 0
    for i, part in enumerate(np.array_split(fix, k)):
        path = os.path.join(d, f"delta-{i:04d}.parquet")
        pq.write_table(pa.table({"id": part, "amount": np.zeros(len(part), dtype=np.int64)}), path); total += os.path.getsize(path)
    return total
c = duckdb.connect(); c.execute("SET threads = 4")
SCAN = f"""SELECT count(*), sum(coalesce(dl.amount, b.amount)) FROM read_parquet('{d}/base-*.parquet') b
           LEFT JOIN read_parquet('{d}/delta-*.parquet') dl USING (id)"""
rows = []
for k in map(int, args.deltas.split(",")):
    mor_bytes = write_deltas(k)
    c.execute(SCAN).fetchall(); ts = []
    for _ in range(3):
        t = time.perf_counter(); c.execute(SCAN).fetchall(); ts.append(time.perf_counter() - t)
    rows.append({"delta_files": k, "delta_mb": mor_bytes / 2**20, "scan_s": min(ts)})
for x in os.listdir(d):
    if x.startswith("delta-"): os.remove(os.path.join(d, x))
t = time.perf_counter(); c.execute(f"SELECT count(*), sum(amount) FROM read_parquet('{d}/base-*.parquet')").fetchall(); plain = time.perf_counter() - t
bm = BitMap(sorted(rng.choice(args.rows_per_file, 1000, replace=False).tolist()))
summary = {"touched_files": len(touched), "files": args.files, "cow_mb": cow_bytes / 2**20, "cow_s": cow_s, "changed_mb": args.corrections * 16 / 2**20,
           "base_mb": base_bytes / 2**20, "plain_scan_s": plain, "roaring_bytes": len(bm.serialize()), "raw_bitmap_bytes": args.rows_per_file // 8}
summary["cow_amplification"] = summary["cow_mb"] / summary["changed_mb"]
shutil.rmtree(d, ignore_errors=True)
print(f"copy-on-write: {len(touched)} of {args.files} files rewritten, {summary['cow_mb']:.0f} MB written for {args.corrections:,} corrections (~{summary['changed_mb']:.2f} MB of change)")
for r in rows: print(f"merge-on-read, {r['delta_files']:>3} delta files ({r['delta_mb']:.2f} MB): full scan {r['scan_s']:.3f} s   (no deltas {plain:.3f} s)")
print(f"Roaring bitmap of 1,000 positions in {args.rows_per_file:,} rows: {summary['roaring_bytes']:,} bytes (raw bitmap {summary['raw_bitmap_bytes']:,})")

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