social-data-attribution-w2/scripts/sampling/split_by_bin.py

"""Split a merged job file into per-bin (topic x format) zst files.

Reads a single job_N.jsonl.zst from --input-dir, routes each document
to the corresponding {topic}__{format}.jsonl.zst file under
--output-dir/job_N/.

Since multiple jobs run concurrently, each writes to its own subdirectory.
Use merge_bin_splits.py afterwards to concatenate into 576 final files.

Usage:
    python scripts/sampling/split_by_bin.py \
        --input-dir stratified_data/merged_docs \
        --output-dir stratified_data/by_bin \
        --job-id 0
"""

import argparse
import io
import json
import logging
from collections import defaultdict
from pathlib import Path

import zstandard

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
)
log = logging.getLogger(__name__)


def split_job(input_path: Path, output_dir: Path) -> dict[str, int]:
    output_dir.mkdir(parents=True, exist_ok=True)

    dctx = zstandard.ZstdDecompressor()
    writers: dict[str, tuple] = {}
    counts: dict[str, int] = defaultdict(int)

    def get_writer(bin_key: str):
        if bin_key not in writers:
            out_path = output_dir / f"{bin_key}.jsonl.zst"
            fh = open(out_path, "wb")
            cctx = zstandard.ZstdCompressor(level=3)
            w = cctx.stream_writer(fh)
            writers[bin_key] = (fh, w)
        return writers[bin_key][1]

    log.info("Reading %s", input_path)
    total = 0

    with open(input_path, "rb") as fin:
        reader = dctx.stream_reader(fin)
        text_in = io.TextIOWrapper(reader, encoding="utf-8")
        for line in text_in:
            doc = json.loads(line)
            topic = doc.get("weborganizer_topic")
            fmt = doc.get("weborganizer_format")
            if not topic or not fmt:
                log.warning("Doc %s missing topic/format, skipping", doc.get("doc_id"))
                continue

            bin_key = f"{topic}__{fmt}"
            w = get_writer(bin_key)
            w.write(line.encode("utf-8") if isinstance(line, str) else line)
            counts[bin_key] += 1
            total += 1

            if total % 500_000 == 0:
                log.info("Processed %d docs (%d bins so far)", total, len(counts))

    for bin_key, (fh, w) in writers.items():
        w.close()
        fh.close()

    log.info("Done: %d docs across %d bins", total, len(counts))
    return dict(counts)


def main() -> None:
    parser = argparse.ArgumentParser(description="Split merged docs into per-bin files")
    parser.add_argument("--input-dir", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--job-id", type=int, required=True, help="Job id (0-4)")
    args = parser.parse_args()

    input_path = args.input_dir / f"job_{args.job_id}.jsonl.zst"
    if not input_path.exists():
        log.error("Input file not found: %s", input_path)
        raise SystemExit(1)

    job_output_dir = args.output_dir / f"job_{args.job_id}"
    counts = split_job(input_path, job_output_dir)

    summary_path = job_output_dir / "bin_counts.json"
    with open(summary_path, "w") as f:
        json.dump(counts, f, indent=2, sort_keys=True)
    log.info("Wrote bin counts to %s", summary_path)


if __name__ == "__main__":
    main()