# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Z-order against sort: count the files.

1 million rows with four uniform columns A–D, clustered into 10,000 files
(a) by sorting on A and (b) by Z-order on A–D (bits interleaved). From each
file's min/max, counts the files a filter on A alone must read, then a filter
on D alone, and compares with N^(1 − k/d).
"""
import argparse, json
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--files", type=int, default=10_000, help="Files.")
p.add_argument("--rows-per-file", type=int, default=100, help="Rows per file.")
p.add_argument("--queries", type=int, default=50, help="Random point filters per column.")
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()

rng = np.random.default_rng(args.seed); n = args.files * args.rows_per_file
cols = rng.integers(0, 2**16, (n, 4)).astype(np.uint64)
def zkey(x):
    z = np.zeros(len(x), dtype=np.uint64)
    for bit in range(16):
        for c in range(4): z |= ((x[:, c] >> np.uint64(bit)) & np.uint64(1)) << np.uint64(bit * 4 + c)
    return z
layouts = {"sort on A": np.argsort(cols[:, 0], kind="stable"), "Z-order on A-D": np.argsort(zkey(cols), kind="stable")}
rows = []
for name, order in layouts.items():
    files = cols[order].reshape(args.files, args.rows_per_file, 4); mn, mx = files.min(axis=1), files.max(axis=1)
    for ci, cname in ((0, "A"), (3, "D")):
        reads = [int(((mn[:, ci] <= v) & (v <= mx[:, ci])).sum()) for v in rng.integers(0, 2**16, args.queries)]
        rows.append({"layout": name, "filter_on": cname, "files_read": float(np.mean(reads)), "files_read_pct": float(np.mean(reads)) / args.files * 100})
summary = {"formula": args.files ** 0.75, "files": args.files}
for r in rows: print(f"{r['layout']:<16} filter on {r['filter_on']}: reads {r['files_read']:8.1f} of {args.files:,} files ({r['files_read_pct']:5.1f}%)")
print(f"N^(1 - 1/4) = {summary['formula']:.0f}")

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