Measure the cost of a boundary: tuple-at-a-time against batches

# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Measure the cost of a boundary: tuple-at-a-time against batches.

The same query three ways over one column: (a) a chain of iterator objects
with one next() call per tuple, (b) the same chain passing batches,
(c) one NumPy expression.
"""
import argparse, json, statistics, time
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=5_000_000, help="Values in the column.")
p.add_argument("--batch-sizes", type=lambda s: [int(v) for v in s.split(",")],
               default=[1, 16, 256, 1024, 16384, 1_000_000], help="Batch sizes to try, comma-separated.")
p.add_argument("--repeats", type=int, default=3, help="Runs per variant; the median time is reported.")
p.add_argument("--seed", type=int, default=0, help="Random seed for the generated column.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

N = args.rows
rng = np.random.default_rng(args.seed)
status = rng.integers(0, 2, N).astype(np.int8)       # 1 = 'OK'
kwh = rng.random(N)


def timed(fn):
    runs = []
    for _ in range(args.repeats):
        t = time.perf_counter(); fn(); runs.append(time.perf_counter() - t)
    return statistics.median(runs)


# (a) tuple-at-a-time: one next() call per tuple
class Scan:
    def __init__(self): self.i = 0
    def next(self):
        if self.i >= N: return None
        r = (status[self.i], kwh[self.i]); self.i += 1; return r

class Filter:
    def __init__(self, c): self.c = c
    def next(self):
        while (r := self.c.next()) is not None:
            if r[0] == 1: return r
        return None

def tuple_at_a_time():
    f = Filter(Scan()); tot = 0.0
    while (r := f.next()) is not None: tot += r[1]
    return tot


# (b) batches: one call per batch
def batches(size):
    for i in range(0, N, size):
        s, k = status[i:i+size], kwh[i:i+size]
        yield k[s == 1]

def batched(size):
    return sum(b.sum() for b in batches(size))


# (c) one expression
def one_expression():
    return kwh[status == 1].sum()


t_tuple = timed(tuple_at_a_time)
print(f"tuple-at-a-time      {t_tuple:8.3f} s")
rows = []
for size in args.batch_sizes:
    t = timed(lambda: batched(size))
    rows.append({"batch_size": size, "seconds": t})
    print(f"batch {size:>12,}   {t:8.3f} s")
t_expr = timed(one_expression)
print(f"one expression       {t_expr:8.3f} s")

by_size = {r["batch_size"]: r["seconds"] for r in rows}
summary = {"tuple_s": t_tuple, "expr_s": t_expr}
if 1 in by_size and 1024 in by_size:
    summary |= {"ratio_vs_1024": t_tuple / by_size[1024], "F_ns": (by_size[1] - by_size[1024]) / N * 1e9}
    print(f"\ntuple-at-a-time / batch 1,024: {summary['ratio_vs_1024']:.0f}x")
    print(f"F (batch 1 minus batch 1,024, per value): {summary['F_ns']:.0f} ns")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {
            "batches": rows,
            "references": [{"variant": "tuple-at-a-time", "seconds": t_tuple},
                           {"variant": "one expression", "seconds": t_expr}],
            "summary": [summary],
        }}, f)