giant/scripts/steps_to_parquet_parallel.py

#!/usr/bin/env python3
"""Convert many ROOT files to Parquet by fanning out to steps_to_parquet.py.

steps_to_parquet.py itself converts a list of files one at a time; this wraps
it to run up to --jobs conversions concurrently, each as its own subprocess
(invoked with the same Python executable running this script, so it picks up
the active venv/uv environment automatically).

Inputs must live under <dataset-root>/raw/<kind>/<gen>/<detector>/<file>.root
(see scripts/migrate_geant_steps.py) — each is written to the matching
processed/<kind>/<gen>/<schema>/<detector>/<file>.parquet, where <schema>
defaults to the highest schemaN already under processed/<kind>/<gen>/ (pass
--schema to pick a specific one, e.g. one just created by
bump_dataset_version.py bump-schema). A file that doesn't fit that layout is
rejected up front, before any conversion runs.

Usage:
    steps_to_parquet_parallel.py raw/steps/gen1/pbwo4/shard-000.root raw/steps/gen1/pbwo4/shard-001.root
    steps_to_parquet_parallel.py raw/steps/gen1/pbwo4/*.root --jobs 8
    steps_to_parquet_parallel.py raw/steps/gen1/pbwo4/*.root --schema schema2
"""

import argparse
import re
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

_STEPS_TO_PARQUET = Path(__file__).resolve().parent / "steps_to_parquet.py"

# Must match scripts/bump_dataset_version.py's GEN_RE / SCHEMA_RE.
GEN_RE = re.compile(r"^gen\d+$")
SCHEMA_RE = re.compile(r"^schema(\d+)$")


class DestinationError(ValueError):
    pass


def latest_schema_tag(processed_gen_dir: Path) -> str | None:
    """Highest schemaN dir directly under *processed_gen_dir*, or None if none exist."""
    if not processed_gen_dir.is_dir():
        return None
    best_tag, best_n = None, -1
    for child in processed_gen_dir.iterdir():
        m = SCHEMA_RE.match(child.name)
        if m and child.is_dir() and int(m.group(1)) > best_n:
            best_tag, best_n = child.name, int(m.group(1))
    return best_tag


def resolve_destination(
    root_file: Path, dataset_root: Path, schema_override: str | None
) -> Path:
    """Map raw/<kind>/<gen>/<detector>/<file>.root (relative to *dataset_root*)
    to processed/<kind>/<gen>/<schema>/<detector>/<file>.parquet.

    Raises DestinationError if *root_file* doesn't fit that layout, or if no
    schema can be determined (no --schema and none exists yet for that gen).
    """
    dataset_root = dataset_root.resolve()
    root_file = root_file.resolve()
    try:
        rel = root_file.relative_to(dataset_root)
    except ValueError:
        raise DestinationError(f"{root_file} is not under dataset root {dataset_root}")

    parts = rel.parts
    if (
        len(parts) != 5
        or parts[0] != "raw"
        or not GEN_RE.match(parts[2])
        or not parts[4].endswith(".root")
    ):
        raise DestinationError(
            f"{root_file} does not match raw/<kind>/<gen>/<detector>/<file>.root "
            f"under {dataset_root} (got relative path: {rel})"
        )
    _, kind, gen_tag, detector, filename = parts
    shard_stem = Path(filename).stem

    processed_gen_dir = dataset_root / "processed" / kind / gen_tag
    schema_tag = schema_override or latest_schema_tag(processed_gen_dir)
    if schema_tag is None:
        raise DestinationError(
            f"no schema exists yet under {processed_gen_dir} — pass --schema or run "
            "bump_dataset_version.py bump-schema first"
        )
    return processed_gen_dir / schema_tag / detector / f"{shard_stem}.parquet"


def _convert_one(
    root_file: str,
    batch_size: str,
    tree: str,
    compression: str,
    steps_to_parquet_path: Path,
    output_path: Path | None,
) -> tuple[str, int, str, str]:
    cmd = [
        sys.executable,
        str(steps_to_parquet_path),
        root_file,
        "--batch-size",
        batch_size,
        "--tree",
        tree,
        "--compression",
        compression,
    ]
    if output_path is not None:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        cmd += ["--output", str(output_path)]
    result = subprocess.run(cmd, capture_output=True, text=True)
    return root_file, result.returncode, result.stdout, result.stderr


def run_parallel(
    root_files: list[str],
    jobs: int,
    batch_size: str = "100 MB",
    tree: str = "Steps",
    compression: str = "snappy",
    steps_to_parquet_path: Path = _STEPS_TO_PARQUET,
    output_for: dict[str, Path] | None = None,
) -> list[tuple[str, int, str, str]]:
    """Run one steps_to_parquet.py subprocess per file, up to *jobs* at a time.

    *output_for*, if given, maps each root_file to the parquet path it should
    be written to (passed through as steps_to_parquet.py's --output); files
    missing from the map fall back to steps_to_parquet.py's own default
    (parquet written next to the input .root).

    Returns one (root_file, returncode, stdout, stderr) tuple per file, in
    completion order (not necessarily input order).
    """
    output_for = output_for or {}
    results = []
    with ThreadPoolExecutor(max_workers=jobs) as pool:
        futures = {
            pool.submit(
                _convert_one,
                f,
                batch_size,
                tree,
                compression,
                steps_to_parquet_path,
                output_for.get(f),
            ): f
            for f in root_files
        }
        for future in as_completed(futures):
            root_file, code, stdout, stderr = future.result()
            status = "ok" if code == 0 else f"FAILED (exit {code})"
            print(f"\n=== {root_file}: {status} ===")
            if stdout:
                print(stdout, end="")
            if stderr:
                print(stderr, end="", file=sys.stderr)
            results.append((root_file, code, stdout, stderr))
    return results


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Convert many ROOT files to Parquet in parallel via steps_to_parquet.py."
    )
    parser.add_argument("root_files", nargs="+", help="Input ROOT file(s)")
    parser.add_argument(
        "-j",
        "--jobs",
        type=int,
        default=4,
        help="Number of conversions to run in parallel (default: 4)",
    )
    parser.add_argument(
        "--batch-size",
        default="100 MB",
        help="Uproot read batch size (default: '100 MB'). E.g. '50 MB', '500000' (rows).",
    )
    parser.add_argument(
        "--tree",
        default="Steps",
        help="Tree name inside the ROOT file (default: Steps)",
    )
    parser.add_argument(
        "--compression",
        default="snappy",
        choices=["snappy", "lz4", "zstd", "gzip", "none"],
        help="Parquet compression codec (default: snappy)",
    )
    parser.add_argument(
        "--dataset-root",
        default="/ceph/lbogner/geant_steps",
        help="Dataset root containing raw/ and processed/ (default: /ceph/lbogner/geant_steps)",
    )
    parser.add_argument(
        "--schema",
        default=None,
        help="Schema tag to write parquets under, e.g. schema2 "
        "(default: highest schemaN already under processed/<kind>/<gen>/)",
    )
    return parser


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    if args.jobs < 1:
        parser.error("--jobs must be >= 1")

    dataset_root = Path(args.dataset_root)
    output_for: dict[str, Path] = {}
    errors: list[str] = []
    for f in args.root_files:
        try:
            output_for[f] = resolve_destination(Path(f), dataset_root, args.schema)
        except DestinationError as exc:
            errors.append(str(exc))

    if errors:
        for err in errors:
            print(f"error: {err}", file=sys.stderr)
        sys.exit(1)

    for root_file, dest in output_for.items():
        print(f"{root_file}  ->  {dest}")

    results = run_parallel(
        args.root_files,
        jobs=args.jobs,
        batch_size=args.batch_size,
        tree=args.tree,
        compression=args.compression,
        output_for=output_for,
    )

    failures = [root_file for root_file, code, _, _ in results if code != 0]
    if failures:
        print(f"\n{len(failures)} of {len(results)} conversion(s) failed:", file=sys.stderr)
        for root_file in failures:
            print(f"  {root_file}", file=sys.stderr)
        sys.exit(1)

    print(f"\nAll {len(results)} conversion(s) completed.")


if __name__ == "__main__":
    main()