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

"""Merge per-job bin splits into final topic x format files.

After split_by_bin.py produces per-job subdirectories, this script
concatenates the zst frames from all jobs into one file per bin.
Zstandard frames are independently decodable, so simple byte
concatenation produces valid multi-frame archives.

Supports sharded execution: pass --worker-id and --num-workers to
process only a slice of the bins (for SLURM array parallelism).

Usage:
    python scripts/sampling/merge_bin_splits.py \
        --input-dir stratified_data/by_bin \
        --output-dir stratified_data/by_bin \
        --jobs 5

    python scripts/sampling/merge_bin_splits.py \
        --input-dir stratified_data/by_bin \
        --output-dir stratified_data/by_bin \
        --jobs 5 \
        --worker-id 0 \
        --num-workers 20
"""

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

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

BUFFER_SIZE = 1024 * 1024


def collect_bins(input_dir: Path, num_jobs: int) -> dict[str, list[Path]]:
    all_bins: dict[str, list[Path]] = defaultdict(list)
    for job_id in range(num_jobs):
        job_dir = input_dir / f"job_{job_id}"
        if not job_dir.exists():
            log.warning("Missing job dir: %s", job_dir)
            continue
        for zst_file in sorted(job_dir.glob("*.jsonl.zst")):
            bin_key = zst_file.stem.replace(".jsonl", "")
            all_bins[bin_key].append(zst_file)
    return all_bins


def collect_counts(
    input_dir: Path, num_jobs: int, bin_keys: list[str]
) -> dict[str, int]:
    total_counts: dict[str, int] = defaultdict(int)
    key_set = set(bin_keys)
    for job_id in range(num_jobs):
        counts_file = input_dir / f"job_{job_id}" / "bin_counts.json"
        if not counts_file.exists():
            continue
        with open(counts_file) as f:
            for k, v in json.load(f).items():
                if k in key_set:
                    total_counts[k] += v
    return dict(total_counts)


def merge_bins(
    all_bins: dict[str, list[Path]],
    output_dir: Path,
) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)

    for bin_key, sources in sorted(all_bins.items()):
        out_path = output_dir / f"{bin_key}.jsonl.zst"
        if len(sources) == 1:
            shutil.copy2(sources[0], out_path)
        else:
            with open(out_path, "wb") as fout:
                for src in sources:
                    with open(src, "rb") as fin:
                        while True:
                            chunk = fin.read(BUFFER_SIZE)
                            if not chunk:
                                break
                            fout.write(chunk)

    log.info("Merged %d bin files into %s", len(all_bins), output_dir)


def main() -> None:
    parser = argparse.ArgumentParser(description="Merge per-job bin splits")
    parser.add_argument("--input-dir", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--jobs", type=int, default=5)
    parser.add_argument(
        "--worker-id", type=int, default=None, help="Worker index (0-based)"
    )
    parser.add_argument("--num-workers", type=int, default=1, help="Total workers")
    args = parser.parse_args()

    all_bins = collect_bins(args.input_dir, args.jobs)
    sorted_keys = sorted(all_bins.keys())
    log.info("Found %d unique bins across %d jobs", len(sorted_keys), args.jobs)

    if args.worker_id is not None:
        my_keys = [
            k
            for i, k in enumerate(sorted_keys)
            if i % args.num_workers == args.worker_id
        ]
        log.info(
            "Worker %d/%d handling %d bins",
            args.worker_id,
            args.num_workers,
            len(my_keys),
        )
    else:
        my_keys = sorted_keys

    my_bins = {k: all_bins[k] for k in my_keys}
    merge_bins(my_bins, args.output_dir)

    my_counts = collect_counts(args.input_dir, args.jobs, my_keys)
    suffix = f"_worker_{args.worker_id}" if args.worker_id is not None else ""
    summary_path = args.output_dir / f"bin_counts{suffix}.json"
    with open(summary_path, "w") as f:
        json.dump(dict(sorted(my_counts.items())), f, indent=2)
    log.info("Wrote bin counts to %s", summary_path)


if __name__ == "__main__":
    main()