# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Encoding by hand: dictionary, RLE and delta.

A 10-million-value country column (200 distinct values, 12 bytes each) and a
sorted timestamp column. Computes the distinct count and bits per value, the
dictionary-encoded size against the raw size, the run count (RLE) before and
after sorting, delta sizes for the timestamps, and gzip of the raw bytes.
"""
import argparse, gzip, json, math
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--n", type=int, default=10_000_000, help="Values.")
p.add_argument("--distinct", type=int, default=200, help="Distinct countries.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

rng = np.random.default_rng(0)
names = np.array([f"country-{i:04d}"[:12].ljust(12) for i in range(args.distinct)], dtype="S12")
codes = rng.zipf(1.3, args.n) % args.distinct; col = names[codes]
raw = col.nbytes; k = len(np.unique(col)); bits = math.ceil(math.log2(k))
dict_size = args.n * bits / 8 + k * 12
runs = lambda a: int((a[1:] != a[:-1]).sum()) + 1
r_unsorted, r_sorted = runs(codes), runs(np.sort(codes))
rle_size = lambda r: r * (bits / 8 + 4)                                    # value code + 4-byte run length
ts = np.sort(rng.integers(1_700_000_000_000_000, 1_700_000_000_000_000 + 86_400_000_000, args.n))   # one day of microseconds
deltas = np.diff(ts); dbits = int(deltas.max()).bit_length()
gz = len(gzip.compress(col[:1_000_000].tobytes(), 6)) * (args.n / 1_000_000)
rows = [{"encoding": "raw strings", "mb": raw / 1e6, "ratio": 1.0},
        {"encoding": f"dictionary ({k} values, {bits} bits each)", "mb": dict_size / 1e6, "ratio": raw / dict_size},
        {"encoding": f"RLE, unsorted ({r_unsorted:,} runs)", "mb": rle_size(r_unsorted) / 1e6, "ratio": raw / rle_size(r_unsorted)},
        {"encoding": f"RLE, sorted ({r_sorted:,} runs)", "mb": rle_size(r_sorted) / 1e6, "ratio": raw / rle_size(r_sorted)},
        {"encoding": "gzip of the raw bytes", "mb": gz / 1e6, "ratio": raw / gz}]
ts_rows = [{"encoding": "raw int64 timestamps", "mb": ts.nbytes / 1e6, "ratio": 1.0},
           {"encoding": f"delta, {dbits} bits per value", "mb": args.n * dbits / 8 / 1e6, "ratio": ts.nbytes / (args.n * dbits / 8)}]
summary = {"distinct": k, "bits": bits, "runs_unsorted": r_unsorted, "runs_sorted": r_sorted, "delta_bits": dbits}
for r in rows + ts_rows: print(f"{r['encoding']:<40} {r['mb']:9.2f} MB  ({r['ratio']:7.1f}x)")

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