social-data-attribution-w2/scripts/transform/combine_merged_docs.py

"""Combine per-worker merged doc files into one per job.

Reads .docs.jsonl.zst files in the input directory for a given job id
and writes one combined job_N.jsonl.zst. When --job-id is omitted,
combines all jobs found.

Usage:
    python scripts/transform/combine_merged_docs.py \
        --input-dir stratified_data/merged_docs \
        --output-dir stratified_data/merged_docs \
        --job-id 0
"""

import argparse
import io
import logging
from pathlib import Path

import zstandard

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


def combine_job(input_dir: Path, output_dir: Path, job_prefix: str) -> None:
    pattern = f"{job_prefix}_worker_*.docs.jsonl.zst"
    files = sorted(input_dir.glob(pattern))
    if not files:
        log.warning("No files matching %s in %s", pattern, input_dir)
        return

    output_path = output_dir / f"{job_prefix}.jsonl.zst"
    log.info("Combining %d files for %s -> %s", len(files), job_prefix, output_path)

    dctx = zstandard.ZstdDecompressor()
    cctx = zstandard.ZstdCompressor(level=3)
    doc_count = 0

    with open(output_path, "wb") as fout:
        writer = cctx.stream_writer(fout)
        for worker_file in files:
            with open(worker_file, "rb") as fin:
                reader = dctx.stream_reader(fin)
                text_in = io.TextIOWrapper(reader, encoding="utf-8")
                for line in text_in:
                    writer.write(line.encode("utf-8"))
                    doc_count += 1
        writer.close()

    log.info("Wrote %d docs to %s", doc_count, output_path)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Combine per-worker docs into per-job 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, default=None, help="Single job to combine (0-4)"
    )
    args = parser.parse_args()

    args.output_dir.mkdir(parents=True, exist_ok=True)

    if args.job_id is not None:
        combine_job(args.input_dir, args.output_dir, f"job_{args.job_id}")
    else:
        for job_id in range(5):
            combine_job(args.input_dir, args.output_dir, f"job_{job_id}")


if __name__ == "__main__":
    main()