# /// script
# requires-python = ">=3.10"
# dependencies = ["duckdb"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Verify a rewrite.

A Parquet dataset of 1 million rows in 50 files is compacted into one sorted
file. count(*) and an order-independent fingerprint, sum(hash(row)), are
compared before and after; then the output is corrupted twice: one row
dropped, and one row replaced by a copy of another (the count stays right).
"""
import argparse, json, os, shutil, tempfile
import duckdb

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=1_000_000, help="Rows.")
p.add_argument("--files", type=int, default=50, help="Input files.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

d = tempfile.mkdtemp(); c = duckdb.connect(); c.execute("SELECT setseed(0.2)")
c.execute(f"CREATE TABLE src AS SELECT i AS id, (random() * 1000)::INT AS amount, 'c' || (i % 977) AS customer FROM range({args.rows}) t(i)")
for i in range(args.files):
    c.execute(f"COPY (SELECT * FROM src WHERE id % {args.files} = {i}) TO '{d}/in-{i:03d}.parquet' (FORMAT parquet)")
FP = "SELECT count(*), sum(hash(id, amount, customer)) FROM read_parquet('{p}')"
before = c.execute(FP.format(p=f"{d}/in-*.parquet")).fetchone()
c.execute(f"COPY (SELECT * FROM read_parquet('{d}/in-*.parquet') ORDER BY customer, id) TO '{d}/compacted.parquet' (FORMAT parquet)")
cases = [("compaction (sorted, one file)", f"SELECT * FROM read_parquet('{d}/compacted.parquet')"),
         ("one row dropped", f"SELECT * FROM read_parquet('{d}/compacted.parquet') WHERE id <> 12345"),
         ("one row replaced by a copy of another", f"SELECT * FROM read_parquet('{d}/compacted.parquet') WHERE id <> 12345 UNION ALL SELECT * FROM read_parquet('{d}/compacted.parquet') WHERE id = 54321")]
rows = []
for name, sql in cases:
    c.execute(f"COPY ({sql}) TO '{d}/out.parquet' (FORMAT parquet)"); after = c.execute(FP.format(p=f"{d}/out.parquet")).fetchone()
    rows.append({"output": name, "count_matches": "yes" if after[0] == before[0] else "no", "fingerprint_matches": "yes" if after[1] == before[1] else "no"})
shutil.rmtree(d, ignore_errors=True)
for r in rows: print(f"{r['output']:<40} count matches: {r['count_matches']:<3}  fingerprint matches: {r['fingerprint_matches']}")

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