giant/scripts/steps_to_parquet_parallel.py

"""Convert many ROOT files to Parquet by fanning out to `dwarf convert`.

A single `dwarf convert` call converts a list of files one at a time; this
module runs up to --jobs conversions concurrently, each as its own `dwarf
convert` subprocess (invoked via `python -m scripts.dwarf`, 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 `dwarf bump-schema`).
A file that doesn't fit that layout is rejected up front, before any
conversion runs.

See `uv run dwarf convert --help` for the CLI.
"""

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

# 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+)$")

# Matches the per-file orphan-drop message printed by
# steps_to_parquet._add_secondary_attributes — each subprocess's count is
# parsed back out of its captured stdout since there's no in-process return
# value across the subprocess boundary.
_ORPHAN_RE = re.compile(r"dropping (\d+) orphaned child_track_id")


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"


_DWARF_CONVERT_CMD = [sys.executable, "-m", "scripts.dwarf", "convert"]


def _convert_one(
    root_file: str,
    batch_size: str,
    tree: str,
    compression: str,
    output_path: Path | None,
    cmd_prefix: list[str],
) -> tuple[str, int, str, str]:
    cmd = [
        *cmd_prefix,
        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",
    output_for: dict[str, Path] | None = None,
    cmd_prefix: list[str] | None = None,
) -> list[tuple[str, int, str, str]]:
    """Run one `dwarf convert` 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 `dwarf convert`'s --output); files
    missing from the map fall back to `dwarf convert`'s own default (parquet
    written next to the input .root).

    *cmd_prefix* overrides the subprocess command run per file (defaults to
    `python -m scripts.dwarf convert`) — used by tests to substitute a fake
    conversion script.

    Returns one (root_file, returncode, stdout, stderr) tuple per file, in
    completion order (not necessarily input order).
    """
    output_for = output_for or {}
    cmd_prefix = cmd_prefix if cmd_prefix is not None else _DWARF_CONVERT_CMD
    results = []
    with ThreadPoolExecutor(max_workers=jobs) as pool:
        futures = {
            pool.submit(
                _convert_one,
                f,
                batch_size,
                tree,
                compression,
                output_for.get(f),
                cmd_prefix,
            ): 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 run_parallel_job(
    root_files: list[str],
    jobs: int,
    dataset_root: Path,
    schema: str | None,
    batch_size: str,
    tree: str,
    compression: str,
) -> None:
    """Resolve each file's dataset-layout destination, convert in parallel, and
    report results. Exits the process (via SystemExit) on destination or
    conversion failure — this is the top-level entry point `dwarf convert`
    delegates to when --jobs > 1."""
    output_for: dict[str, Path] = {}
    errors: list[str] = []
    for f in root_files:
        try:
            output_for[f] = resolve_destination(Path(f), dataset_root, schema)
        except DestinationError as exc:
            errors.append(str(exc))

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

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

    results = run_parallel(
        root_files,
        jobs=jobs,
        batch_size=batch_size,
        tree=tree,
        compression=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)
        raise SystemExit(1)

    total_orphaned = sum(
        int(m.group(1)) for _, _, stdout, _ in results for m in _ORPHAN_RE.finditer(stdout)
    )
    if total_orphaned:
        print(f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s).")
    print(f"\nAll {len(results)} conversion(s) completed.")