# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "pandas", "pyarrow"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Per-row UDF against vectorized UDF.

Scores 5 million rows with x * 0.3 + 2.0: (a) a Python function called once
per row (pandas apply), (b) a vectorized function over blocks of 100, 10,000
and 1,000,000 rows. Also times reading the input from Parquet and writing the
result without the function.
"""
import argparse, json, os, tempfile, time
import numpy as np, pandas as pd

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=5_000_000, help="Rows.")
p.add_argument("--blocks", default="100,10000,1000000", help="Block sizes for the vectorized function.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

d = tempfile.mkdtemp(); src = os.path.join(d, "in.parquet"); out = os.path.join(d, "out.parquet")
pd.DataFrame({"x": np.random.default_rng(0).random(args.rows)}).to_parquet(src)
t = time.perf_counter(); df = pd.read_parquet(src); read_s = time.perf_counter() - t
t = time.perf_counter(); df.assign(score=df.x).to_parquet(out); write_s = time.perf_counter() - t
def per_row(v): return v * 0.3 + 2.0
def vectorized(a): return a * 0.3 + 2.0
rows = []
t = time.perf_counter(); r1 = df.x.apply(per_row); rows.append({"method": "(a) per-row Python function", "block": 1, "seconds": time.perf_counter() - t})
x = df.x.to_numpy()
for b in map(int, args.blocks.split(",")):
    t = time.perf_counter(); r2 = np.concatenate([vectorized(x[i:i + b]) for i in range(0, len(x), b)])
    rows.append({"method": f"(b) vectorized, blocks of {b:,}", "block": b, "seconds": time.perf_counter() - t})
assert np.allclose(r1.to_numpy(), r2)
for f in (src, out): os.remove(f)
base = next(r["seconds"] for r in rows if r["block"] == 10_000)
for r in rows: r["vs_10k_blocks"] = r["seconds"] / base
summary = {"read_s": read_s, "write_s": write_s, "per_row_over_10k": rows[0]["seconds"] / base}
for r in rows: print(f"{r['method']:<34} {r['seconds']:8.3f} s  ({r['vs_10k_blocks']:7.1f}x blocks of 10,000)")
print(f"reading the input {read_s:.3f} s, writing the result {write_s:.3f} s")

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