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

#!/usr/bin/env python3
"""Shard a large JSONL.zst file into smaller JSONL.zst shards.

Single-pass streaming: decompress → count lines → recompress into shards.
Never holds more than one line in memory at a time.

Usage:
    python scripts/transform/shard_cleaned_pool.py \
        --input ~/scratch/archive-dolma3-pool-150b/archive-dolma3-pool-150b-cleaned.jsonl.zst \
        --output-dir /tmp/shards \
        --lines-per-shard 10000000
"""

from __future__ import annotations

import argparse
import io
import sys
import time
from pathlib import Path


def shard_jsonl_zst(
    input_path: Path,
    output_dir: Path,
    lines_per_shard: int,
    compression_level: int = 3,
) -> list[Path]:
    """Stream-shard a .jsonl.zst file into multiple .jsonl.zst shards."""
    import zstandard as zstd

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

    # First pass: count total lines to determine shard naming
    print(f"Pass 1/2: Counting lines in {input_path} ...")
    t0 = time.time()
    total_lines = 0
    dctx = zstd.ZstdDecompressor()
    with input_path.open("rb") as fh:
        with dctx.stream_reader(fh) as reader:
            for line in io.TextIOWrapper(reader, encoding="utf-8"):
                if line.strip():
                    total_lines += 1
                if total_lines % 5_000_000 == 0:
                    print(f"  ... counted {total_lines:,} lines")

    num_shards = (total_lines + lines_per_shard - 1) // lines_per_shard
    elapsed = time.time() - t0
    print(
        f"  Total: {total_lines:,} lines → {num_shards} shards "
        f"({lines_per_shard:,} lines each, last shard may be smaller)"
    )
    print(f"  Count took {elapsed:.0f}s")

    # Second pass: write shards
    print(f"\nPass 2/2: Writing shards to {output_dir} ...")
    t0 = time.time()

    shard_paths: list[Path] = []
    current_shard = 0
    current_line = 0
    cctx = zstd.ZstdCompressor(level=compression_level)
    writer = None
    out_fh = None

    def open_shard(idx: int):
        name = f"train-{idx:05d}-of-{num_shards:05d}.jsonl.zst"
        path = output_dir / name
        shard_paths.append(path)
        fh = path.open("wb")
        w = cctx.stream_writer(fh)
        print(f"  Writing shard {idx + 1}/{num_shards}: {name}")
        return fh, w

    dctx2 = zstd.ZstdDecompressor()
    with input_path.open("rb") as fh:
        with dctx2.stream_reader(fh) as reader:
            out_fh, writer = open_shard(current_shard)
            lines_in_shard = 0

            for line in io.TextIOWrapper(reader, encoding="utf-8"):
                if not line.strip():
                    continue

                # Ensure line ends with \n
                if not line.endswith("\n"):
                    line = line + "\n"

                writer.write(line.encode("utf-8"))
                lines_in_shard += 1
                current_line += 1

                if lines_in_shard >= lines_per_shard and current_shard < num_shards - 1:
                    writer.close()
                    out_fh.close()
                    print(
                        f"    → {lines_in_shard:,} lines, "
                        f"size: {shard_paths[-1].stat().st_size / 1e9:.2f} GB"
                    )

                    current_shard += 1
                    lines_in_shard = 0
                    out_fh, writer = open_shard(current_shard)

                if current_line % 5_000_000 == 0:
                    print(f"  ... processed {current_line:,}/{total_lines:,} lines")

    # Close the last shard
    if writer:
        writer.close()
    if out_fh:
        out_fh.close()
    if shard_paths:
        print(
            f"    → {lines_in_shard:,} lines, "
            f"size: {shard_paths[-1].stat().st_size / 1e9:.2f} GB"
        )

    elapsed = time.time() - t0
    print(f"\nSharding complete in {elapsed / 3600:.1f} hours ({elapsed:.0f}s)")
    print(f"Total: {current_line:,} lines across {len(shard_paths)} shards")

    # Summary
    total_size = sum(p.stat().st_size for p in shard_paths)
    print("\nShard summary:")
    for p in shard_paths:
        sz = p.stat().st_size
        print(f"  {p.name}: {sz / 1e9:.2f} GB")
    print(f"  Total: {total_size / 1e9:.2f} GB")

    return shard_paths


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Shard a large JSONL.zst file")
    parser.add_argument(
        "--input", type=Path, required=True, help="Path to the input .jsonl.zst file"
    )
    parser.add_argument(
        "--output-dir", type=Path, required=True, help="Directory to write shards to"
    )
    parser.add_argument(
        "--lines-per-shard",
        type=int,
        default=10_000_000,
        help="Number of JSONL lines per shard (default: 10M)",
    )
    parser.add_argument(
        "--compression-level",
        type=int,
        default=3,
        help="Zstandard compression level (default: 3)",
    )
    args = parser.parse_args(argv)

    if not args.input.exists():
        print(f"ERROR: Input file not found: {args.input}")
        return 1

    shard_jsonl_zst(
        input_path=args.input,
        output_dir=args.output_dir,
        lines_per_shard=args.lines_per_shard,
        compression_level=args.compression_level,
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())