# /// script
# requires-python = ">=3.10"
# dependencies = ["duckdb"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Rename is drop plus add, in a table.

Two Parquet files, an older one with a column weight and a newer one where it
is called weight_kg, read together by DuckDB with union_by_name; then the
read-side workaround COALESCE(weight_kg, weight).
"""
import argparse, json, os, tempfile
import duckdb

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
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(f"COPY (SELECT 'S1' AS id, 12.5 AS weight) TO '{d}/old.parquet' (FORMAT parquet)")
c.execute(f"COPY (SELECT 'S2' AS id, 3.2 AS weight_kg) TO '{d}/new.parquet' (FORMAT parquet)")
files = f"['{d}/old.parquet', '{d}/new.parquet']"
union = c.execute(f"SELECT id, weight, weight_kg FROM read_parquet({files}, union_by_name = true) ORDER BY id").fetchall()
cols = [r[0] for r in c.execute(f"DESCRIBE SELECT * FROM read_parquet({files}, union_by_name = true)").fetchall()]
fixed = c.execute(f"SELECT id, COALESCE(weight_kg, weight) AS weight_kg FROM read_parquet({files}, union_by_name = true) ORDER BY id").fetchall()
total_naive = c.execute(f"SELECT sum(weight_kg) FROM read_parquet({files}, union_by_name = true)").fetchone()[0]
rows = [{"id": i, "weight": w, "weight_kg": wk, "coalesced": f} for (i, w, wk), (_, f) in zip(union, fixed)]
summary = {"columns": ", ".join(cols), "sum_weight_kg_naive": float(total_naive), "sum_coalesced": float(sum(f for _, f in fixed))}
print("columns:", summary["columns"])
for r in rows: print(f"{r['id']}: weight {r['weight']}, weight_kg {r['weight_kg']}, COALESCE(weight_kg, weight) {r['coalesced']}")
print(f"sum(weight_kg) without the workaround: {total_naive}; with it: {summary['sum_coalesced']}")
for f in os.listdir(d): os.remove(os.path.join(d, f))

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"rows": rows}}, f, default=float)