social-data-attribution-w2/scripts/r2_sunset/sample_inventory.py

"""Show representative keys + size distribution per prefix from the local inventory.

Read-only against runs/r2_sunset/r2_inventory.jsonl. No R2 calls. No cost.

For each prefix: first 5 keys (sorted), last 5, 10 random middle, size stats,
file-extension breakdown, common path-pattern templates.
"""

from __future__ import annotations

import argparse
import json
import random
import re
import statistics
from collections import Counter, defaultdict
from pathlib import Path


def fmt_bytes(b: int) -> str:
    v: float = float(b)
    for u in ("B", "KB", "MB", "GB"):
        if v < 1024:
            return f"{v:.1f} {u}"
        v /= 1024
    return f"{v:.1f} TB"


def extension(key: str) -> str:
    name = key.rsplit("/", 1)[-1]
    if name.endswith(".jsonl.zst"):
        return ".jsonl.zst"
    if name.endswith(".parquet.zst"):
        return ".parquet.zst"
    if "." in name:
        return "." + name.rsplit(".", 1)[1]
    return "(no-ext)"


def path_template(key: str) -> str:
    """Normalize digits to N and uuid-like segments to <id> for pattern recognition."""
    parts = []
    for seg in key.split("/"):
        s = re.sub(r"\d+", "N", seg)
        s = re.sub(r"[0-9a-f]{8,}", "<hex>", s)
        parts.append(s)
    return "/".join(parts)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--inventory", type=Path, default=Path("runs/r2_sunset/r2_inventory.jsonl")
    )
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    rng = random.Random(args.seed)
    keys_by_prefix: dict[str, list[tuple[str, int]]] = defaultdict(list)

    with args.inventory.open() as f:
        for line in f:
            row = json.loads(line)
            keys_by_prefix[row["prefix"]].append((row["key"], row["size"]))

    for prefix in sorted(keys_by_prefix):
        rows = keys_by_prefix[prefix]
        sizes = [s for _, s in rows]
        keys = sorted(k for k, _ in rows)
        n = len(keys)
        print("=" * 80)
        print(f"PREFIX: {prefix}")
        print(f"  count: {n:,}   total: {fmt_bytes(sum(sizes))}")
        print(
            f"  size  min={fmt_bytes(min(sizes))}  median={fmt_bytes(int(statistics.median(sizes)))}  "
            f"mean={fmt_bytes(int(statistics.mean(sizes)))}  max={fmt_bytes(max(sizes))}"
        )

        ext_counter = Counter(extension(k) for k, _ in rows)
        print(f"  extensions: {dict(ext_counter.most_common(8))}")

        tmpl_counter = Counter(path_template(k) for k, _ in rows)
        print("  top 5 path templates:")
        for tmpl, c in tmpl_counter.most_common(5):
            print(f"    {c:>10,}  {tmpl}")

        print("  first 5 keys (lex sorted):")
        for k in keys[:5]:
            print(f"    {k}")
        print("  last 5 keys:")
        for k in keys[-5:]:
            print(f"    {k}")
        sample_n = min(10, n)
        print(f"  {sample_n} random middle keys:")
        for k in rng.sample(keys, sample_n):
            print(f"    {k}")
        print()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())