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

The same data (about 400 MB of synthetic events) written as Parquet in three
layouts: 800 files of 0.5 MB, 8 files of 50 MB, and 1 file of 400 MB. A full
scan of each with DuckDB, from local disk and from an S3-compatible object
store (SeaweedFS), then F (fixed cost per file) and b (bandwidth) fitted from
time = files × F + bytes / b.
Needs the S3 service from the compose file (port 59100) for the S3 half.
"""
import argparse, json, os, shutil, tempfile, time
import boto3, duckdb, numpy as np, pyarrow as pa, pyarrow.parquet as pq
from botocore.config import Config

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=25_000_000, help="Rows in all (about 16 bytes each, compressed).")
p.add_argument("--layouts", default="800,8,1", help="File counts.")
p.add_argument("--endpoint", default="localhost:59100", help="S3 endpoint (host:port).")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

rng = np.random.default_rng(0); work = tempfile.mkdtemp()
tbl = pa.table({"id": np.arange(args.rows), "amount": rng.integers(1, 10_000, args.rows), "ts": rng.integers(0, 10**9, args.rows)})
s3 = boto3.client("s3", endpoint_url=f"http://{args.endpoint}", aws_access_key_id="lab", aws_secret_access_key="lablablab",
                  config=Config(s3={"addressing_style": "path"}), region_name="us-east-1")
try: s3.create_bucket(Bucket="lake")
except Exception: pass
counts = [int(x) for x in args.layouts.split(",")]
for n in counts:
    d = os.path.join(work, f"f{n}"); os.makedirs(d); step = args.rows // n
    for i in range(n):
        path = os.path.join(d, f"part-{i:05d}.parquet"); pq.write_table(tbl.slice(i * step, step), path)
        s3.upload_file(path, "lake", f"f{n}/part-{i:05d}.parquet")
c = duckdb.connect(); c.execute("INSTALL httpfs"); c.execute("LOAD httpfs")
c.execute(f"SET s3_endpoint='{args.endpoint}'"); c.execute("SET s3_use_ssl=false"); c.execute("SET s3_url_style='path'")
c.execute("SET s3_access_key_id='lab'"); c.execute("SET s3_secret_access_key='lablablab'"); c.execute("SET s3_region='us-east-1'")
for setting in ("enable_external_file_cache", "enable_http_metadata_cache", "parquet_metadata_cache"):     # every scan must go to the store
    try: c.execute(f"SET {setting} = false")
    except duckdb.Error: pass
rows = []
for where in ("local disk", "S3 (SeaweedFS)"):
    for n in counts:
        src = os.path.join(work, f"f{n}", "*.parquet") if where == "local disk" else f"s3://lake/f{n}/*.parquet"
        q = f"SELECT count(*), sum(id), sum(amount), sum(ts) FROM read_parquet('{src}')"            # every column is read
        c.execute(q).fetchall(); ts = []
        for _ in range(3):
            t = time.perf_counter(); c.execute(q).fetchall(); ts.append(time.perf_counter() - t)
        mb = sum(os.path.getsize(os.path.join(work, f"f{n}", f)) for f in os.listdir(os.path.join(work, f"f{n}"))) / 2**20
        rows.append({"storage": where, "files": n, "file_mb": mb / n, "total_mb": mb, "seconds": min(ts)})
fits = []
for where in ("local disk", "S3 (SeaweedFS)"):
    pts = [r for r in rows if r["storage"] == where]
    X = np.array([[r["files"], r["total_mb"]] for r in pts]); y = np.array([r["seconds"] for r in pts])
    (F, inv_b), *_ = np.linalg.lstsq(X, y, rcond=None)
    for r in pts:
        s = r["file_mb"]; r["fit_overhead_pct"] = F / (F + s * inv_b) * 100 if F > 0 else 0.0
    fits.append({"storage": where, "F_ms": F * 1000, "b_mb_s": 1 / inv_b if inv_b > 0 else None})
shutil.rmtree(work, ignore_errors=True)
for r in rows: print(f"{r['storage']:<15} {r['files']:>4} files of {r['file_mb']:7.1f} MB: {r['seconds']:6.3f} s   per-file overhead {r['fit_overhead_pct']:5.1f}%")
for f_ in fits: print(f"{f_['storage']:<15} fit: F = {f_['F_ms']:.1f} ms per file, b = {f_['b_mb_s'] or 0:,.0f} MB/s")

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