# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Batch size: find the knee.

200,000 rows inserted into SQLite on disk with one commit per batch, for
batch sizes from 1 to 10,000. Then trickle traffic at 20 rows/s: a
size-only batch of 1,000 against "size OR age >= 200 ms".
"""
import argparse, json, os, sqlite3, tempfile, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=200_000, help="Rows per batch size (batch size 1 inserts 20,000).")
p.add_argument("--sizes", type=lambda s: [int(v) for v in s.split(",")], default=[1, 10, 100, 1000, 10000], help="Batch sizes.")
p.add_argument("--trickle", type=float, default=20, help="Rows per second of trickle traffic.")
p.add_argument("--max-age-ms", type=float, default=200, help="Age limit of the size-or-age policy.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
path = os.path.join(tempfile.gettempdir(), "batch-knee.db")
import atexit; atexit.register(lambda: os.path.exists(path) and os.remove(path))   # leave nothing behind

def bench(batch, n):
    db = sqlite3.connect(path); db.execute("DROP TABLE IF EXISTS t"); db.execute("CREATE TABLE t(id INTEGER, v TEXT)"); db.commit()
    t0 = time.perf_counter()
    for i in range(0, n, batch):
        db.executemany("INSERT INTO t VALUES (?,?)", [(j, "x" * 50) for j in range(i, i + batch)]); db.commit()
    dt = time.perf_counter() - t0; db.close()
    return n / dt

rows = []
for b in args.sizes:
    r = bench(b, 20_000 if b == 1 else args.rows)
    rows.append({"batch": b, "rows_per_s": r}); print(f"batch {b:>6,}: {r:>10,.0f} rows/s")

# trickle traffic, simulated clock: how long does the last row of a batch wait?
arrivals = [i / args.trickle for i in range(1000)]
size_only = 999 / args.trickle                                 # the first row waits for 999 more
waits, start = [], arrivals[0]
for t in arrivals:
    if t - start >= args.max_age_ms / 1000: start = t          # the age limit flushed the previous batch
    waits.append(t - start)
size_or_age = max(waits) + 1 / args.trickle
print(f"trickle at {args.trickle:g} rows/s: size-only batch of 1,000 makes its first row wait {size_only:.1f} s; "
      f"size OR age >= {args.max_age_ms:g} ms bounds every wait at {size_or_age * 1000:.0f} ms")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"batches": rows, "summary": [{"size_only_wait_s": size_only, "size_or_age_wait_ms": size_or_age * 1000,
                   "knee_ratio": rows[-1]["rows_per_s"] / rows[0]["rows_per_s"]}]}}, f)