# /// script
# requires-python = ">=3.10"
# dependencies = ["duckdb", "numpy", "pyarrow"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Find the compaction interval.

Measures two constants on a local Parquet table: c, the extra scan time per
unmerged delta file, and bU, the time to rewrite the table (compaction).
With a delta arrival rate w and a read rate a, minimises bU/τ + a·c·w·τ/2 at τ* = √(2·bU / (a·c·w))
and then simulates a day of the workload compacting every τ*/4, τ* and 4τ*,
adding up read overhead and compaction time.
"""
import argparse, json, math, os, shutil, tempfile, time
import duckdb, numpy as np, pyarrow as pa, pyarrow.parquet as pq

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=5_000_000, help="Rows in the table.")
p.add_argument("--w", type=float, default=60, help="Delta files arriving per hour.")
p.add_argument("--a", type=float, default=120, help="Full scans per hour.")
p.add_argument("--hours", type=float, default=24, help="Simulated hours.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

rng = np.random.default_rng(0); d = tempfile.mkdtemp()
pq.write_table(pa.table({"id": np.arange(args.rows), "amount": rng.integers(1, 10_000, args.rows)}), os.path.join(d, "base.parquet"))
c = duckdb.connect(); c.execute("SET threads = 4")
def scan_time(k):
    for x in [x for x in os.listdir(d) if x.startswith("delta-")]: os.remove(os.path.join(d, x))
    for i in range(k):
        ids = rng.choice(args.rows, 50, replace=False); pq.write_table(pa.table({"id": ids, "amount": np.zeros(50, dtype=np.int64)}), os.path.join(d, f"delta-{i:04d}.parquet"))
    src = f"read_parquet('{d}/delta-*.parquet')" if k else "(SELECT 0::BIGINT AS id, 0::BIGINT AS amount WHERE false)"
    q = f"SELECT sum(coalesce(dl.amount, b.amount)) FROM read_parquet('{d}/base.parquet') b LEFT JOIN {src} dl USING (id)"
    c.execute(q).fetchall(); ts = []
    for _ in range(3):
        t = time.perf_counter(); c.execute(q).fetchall(); ts.append(time.perf_counter() - t)
    return min(ts)
ks = [0, 10, 40, 80]; times = [scan_time(k) for k in ks]
cost_per_delta = float(np.polyfit(ks, times, 1)[0])                     # c: seconds of extra scan per unmerged delta
t = time.perf_counter(); c.execute(f"COPY (SELECT * FROM read_parquet('{d}/base.parquet')) TO '{d}/rewritten.parquet' (FORMAT parquet)"); bU = time.perf_counter() - t
shutil.rmtree(d, ignore_errors=True)
tau_star = math.sqrt(2 * bU / (args.a * cost_per_delta * args.w))       # hours; cost/h = bU/τ + a·c·w·τ/2
def day(tau):
    read = args.a * cost_per_delta * args.w * tau / 2 * args.hours            # average unmerged deltas = w·τ/2
    compact = bU * args.hours / tau
    return read, compact
rows = []
for label, tau in (("τ*/4", tau_star / 4), ("τ*", tau_star), ("4τ*", 4 * tau_star)):
    r, cpt = day(tau); rows.append({"interval": label, "tau_h": tau, "read_overhead_s": r, "compaction_s": cpt, "total_s": r + cpt})
summary = {"c_ms": cost_per_delta * 1000, "bU_s": bU, "tau_star_h": tau_star, "w": args.w, "a": args.a}
print(f"measured: c = {summary['c_ms']:.2f} ms of scan per unmerged delta file, bU = {bU:.2f} s to rewrite the table; τ* = {tau_star:.2f} h")
for r in rows: print(f"compact every {r['interval']:<5} ({r['tau_h']:5.2f} h): read overhead {r['read_overhead_s']:7.1f} s + compaction {r['compaction_s']:7.1f} s = {r['total_s']:7.1f} s per day")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"intervals": rows, "scans": [{"deltas": k, "scan_s": tm} for k, tm in zip(ks, times)]}}, f)