Overlap depth before and after a blind compaction

# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Overlap depth before and after a compaction.

200 small files written in arrival order, each holding a narrow range of
customer_id that wanders over time. Overlap depth (files whose [min, max]
contains a point) at 1,000 random points, then after compacting (a) the oldest
files together in time-adjacent groups of 10 and (b) files grouped by key
range into the same number of outputs.
"""
import argparse, json
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--files", type=int, default=200, help="Small files.")
p.add_argument("--group", type=int, default=10, help="Files merged per compaction output.")
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)
centers = rng.uniform(0, 1_000_000, args.files)                         # each file's customers cluster somewhere: arrival order is not key order
files = [np.clip(rng.normal(c, 5_000, 1000), 0, 1_000_000) for c in centers]
def depth(fs):
    mn = np.array([f.min() for f in fs]); mx = np.array([f.max() for f in fs]); pts = rng.uniform(0, 1_000_000, 1000)
    d = np.array([((mn <= x) & (x <= mx)).sum() for x in pts]); return float(d.mean()), int(d.max()), len(fs)
by_time = [np.concatenate(files[i:i + args.group]) for i in range(0, args.files, args.group)]
allv = np.sort(np.concatenate(files)); k = len(by_time)
by_key = np.array_split(allv, k)
rows = []
for name, fs in (("before: small files in arrival order", files), ("(a) compacted by time (oldest together)", by_time), ("(b) compacted by key range", by_key)):
    mean, mx, n = depth(fs); rows.append({"layout": name, "files": n, "mean_depth": mean, "max_depth": mx})
for r in rows: print(f"{r['layout']:<42} {r['files']:>4} files   overlap depth mean {r['mean_depth']:6.2f}, max {r['max_depth']}")

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