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

import argparse
import json
import sys
from pathlib import Path

import pyarrow.parquet as pq


def convert_parquet_to_jsonl(
    parquet_path: Path, jsonl_path: Path, batch_size: int = 50_000
):
    pf = pq.ParquetFile(parquet_path)
    row_count = 0
    with open(jsonl_path, "w") as f:
        for batch in pf.iter_batches(batch_size=batch_size):
            rows = batch.to_pydict()
            n = len(next(iter(rows.values())))
            for i in range(n):
                record = {k: rows[k][i] for k in rows}
                f.write(json.dumps(record, default=str) + "\n")
            row_count += n
            print(f"  {row_count} rows written", flush=True)
    print(f"  Total: {row_count} rows -> {jsonl_path}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--index", type=int, required=True, help="Parquet shard index (0-4)"
    )
    parser.add_argument("--data-dir", type=Path, default=Path("stratified_data"))
    parser.add_argument("--total-shards", type=int, default=5)
    args = parser.parse_args()

    pq_file = (
        args.data_dir / f"train-{args.index:05d}-of-{args.total_shards:05d}.parquet"
    )
    jsonl_file = (
        args.data_dir / f"train-{args.index:05d}-of-{args.total_shards:05d}.jsonl"
    )

    if not pq_file.exists():
        print(f"Missing: {pq_file}", file=sys.stderr)
        sys.exit(1)

    print(f"Converting {pq_file}")
    convert_parquet_to_jsonl(pq_file, jsonl_file)
    print("Done.")


if __name__ == "__main__":
    main()