# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Splittability: CSV with quoted newlines against NDJSON.

The same 100,000 records, some with a newline inside a note, as CSV and as
NDJSON. Parallel readers start at random byte offsets, skip to the next
newline and parse: how many records come out broken?
"""
import argparse, csv, json, os, random, tempfile

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--records", type=int, default=100_000, help="Records.")
p.add_argument("--splits", type=int, default=200, help="Random start offsets tried.")
p.add_argument("--chunk", type=int, default=200_000, help="Bytes each split reads.")
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)
d = tempfile.mkdtemp()
rows = [{"id": i, "note": "ok" if i % 7 else "Fragile\nhandle with care"} for i in range(args.records)]
csvp, js = os.path.join(d, "data.csv"), os.path.join(d, "data.ndjson")
with open(csvp, "w", newline="") as f:
    w = csv.DictWriter(f, fieldnames=["id", "note"]); w.writeheader(); w.writerows(rows)
with open(js, "w") as f:
    for r in rows: f.write(json.dumps(r) + "\n")

def split_read(path, start):
    with open(path, "rb") as f:
        f.seek(start); f.readline()
        text = f.read(args.chunk).decode(errors="ignore")
    return text.splitlines()[:-1]                                     # drop the line cut by the read limit
def bad_csv(lines):
    """Parse from the split point with a real CSV reader; a record is broken if its id is not a number
    or it does not have exactly two fields (the split started inside a quoted note)."""
    return sum(1 for rec in csv.reader(lines) if len(rec) != 2 or not rec[0].isdigit())
def bad_json(lines):
    bad = 0
    for ln in lines:
        try: json.loads(ln)
        except json.JSONDecodeError: bad += 1
    return bad
res = {"CSV": [0, 0], "NDJSON": [0, 0]}
for _ in range(args.splits):
    for fmt, path, fn in (("CSV", csvp, bad_csv), ("NDJSON", js, bad_json)):
        s = random.randrange(1000, os.path.getsize(path) - args.chunk)
        b = fn(split_read(path, s)); res[fmt][0] += b; res[fmt][1] += b > 0
out = [{"format": k, "broken_lines": v[0], "splits_with_broken": v[1], "splits": args.splits} for k, v in res.items()]
for r in out: print(f"{r['format']:<7} {r['splits_with_broken']:>3} of {args.splits} splits hit broken records ({r['broken_lines']:,} broken lines)")

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