Normalized keys and sorting a pointer array

# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Normalized keys, and sorting a pointer.

300,000 three-column rows sorted with a typed comparison function, and by
normalized byte-string keys whose byte order equals the tuple order. Then
sorts (8-byte prefix, index) pairs instead of whole rows, compares the memory
of the sort inputs, and checks that all orders agree.
"""
import argparse, functools, json, random, struct, sys, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--n", type=int, default=300_000, help="Rows.")
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()
random.seed(args.seed)

rows = [(random.choice(["north", "south", "east", "west"]), random.randrange(10**6), random.random()) for _ in range(args.n)]
def cmp(a, b):                                                          # typed, branchy comparison
    if a[0] != b[0]: return -1 if a[0] < b[0] else 1
    if a[1] != b[1]: return -1 if a[1] < b[1] else 1
    return (a[2] > b[2]) - (a[2] < b[2])
def norm(r):                                                            # bytes whose order equals the tuple order
    return r[0].encode().ljust(8, b"\0") + struct.pack(">I", r[1]) + struct.pack(">d", r[2])
t = time.perf_counter(); a = sorted(rows, key=functools.cmp_to_key(cmp)); t1 = time.perf_counter() - t
t = time.perf_counter(); keys = sorted((norm(r), i) for i, r in enumerate(rows)); t2 = time.perf_counter() - t
prefixes = [(norm(r)[:8] + struct.pack(">I", r[1]), i) for i, r in enumerate(rows)]
t = time.perf_counter(); pref = sorted(prefixes); t3 = time.perf_counter() - t
size = lambda xs, inner: sys.getsizeof(xs) + sum(sys.getsizeof(x) + sum(sys.getsizeof(y) for y in inner(x)) for x in xs)
mem_rows = size(rows, lambda r: r); mem_pref = size(prefixes, lambda r: r)
same = [rows[i] for _, i in keys] == a
results = [{"method": "typed comparison", "seconds": t1}, {"method": "normalized keys", "seconds": t2},
           {"method": "prefix + index pairs", "seconds": t3}]
summary = {"ratio": t1 / t2, "same_order": "yes" if same else "no", "rows_mb": mem_rows / 2**20, "prefix_mb": mem_pref / 2**20}
for r in results: print(f"{r['method']:<22} {r['seconds']:.3f} s")
print(f"typed / normalized = {summary['ratio']:.1f}x; same order: {summary['same_order']}; sort input: rows {summary['rows_mb']:.0f} MB, prefix pairs {summary['prefix_mb']:.0f} MB")

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