#!/usr/bin/env python3 """Convert the Steps tree from a ROOT file to Parquet. Usage: uv run python steps_to_parquet.py input.root uv run python steps_to_parquet.py input.root -o output.parquet uv run python steps_to_parquet.py input.root --batch-size "200 MB" --tree Hits uv run python steps_to_parquet.py input1.root input2.root input3.root """ import argparse from pathlib import Path from typing import Literal import awkward as ak import polars as pl import uproot ParquetCompression = Literal["lz4", "uncompressed", "snappy", "gzip", "brotli", "zstd"] def _add_secondary_attributes(df: pl.DataFrame) -> pl.DataFrame: """Add per-step secondary attributes via the parent→child track join. For each step that spawns secondaries, collects each child track's birth state (from the child track's first step in the same event) and emits: e_sec float64 — total secondary energy (sum of child first-step pre_E) sec_E_list list[f64] — per-secondary energy, sorted descending sec_pdg_list list[i32] — per-secondary PDG code, same order sec_dx_list list[f64] — per-secondary birth direction x, same order sec_dy_list list[f64] — per-secondary birth direction y, same order sec_dz_list list[f64] — per-secondary birth direction z, same order Steps with no children get 0.0 / empty lists. The full event must be present in `df` (it is — the writer concatenates before calling this). """ first_step = ( df.sort("step_no") .group_by(["event_id", "track_id"]) .agg( pl.col("pre_E").first().alias("child_E"), pl.col("pdg").first().alias("child_pdg"), pl.col("pre_dx").first().alias("child_dx"), pl.col("pre_dy").first().alias("child_dy"), pl.col("pre_dz").first().alias("child_dz"), ) .rename({"track_id": "child_track_id"}) ) exploded = ( df.select(["event_id", "child_track_ids"]) .with_row_index("_step_row") .explode("child_track_ids") .rename({"child_track_ids": "child_track_id"}) .drop_nulls("child_track_id") ) joined = exploded.join(first_step, on=["event_id", "child_track_id"], how="left") # Sort each step's secondaries by descending energy, then aggregate into lists per_step = ( joined.sort("child_E", descending=True) .group_by("_step_row") .agg( pl.col("child_E").sum().alias("e_sec"), pl.col("child_E").alias("sec_E_list"), pl.col("child_pdg").alias("sec_pdg_list"), pl.col("child_dx").alias("sec_dx_list"), pl.col("child_dy").alias("sec_dy_list"), pl.col("child_dz").alias("sec_dz_list"), ) ) empty_list_f64 = pl.Series("x", [[]], dtype=pl.List(pl.Float64)) empty_list_i32 = pl.Series("x", [[]], dtype=pl.List(pl.Int32)) return ( df.with_row_index("_step_row") .join(per_step, on="_step_row", how="left") .with_columns( pl.col("e_sec").fill_null(0.0).cast(pl.Float64), pl.col("sec_E_list").fill_null(empty_list_f64), pl.col("sec_pdg_list").fill_null(empty_list_i32), pl.col("sec_dx_list").fill_null(empty_list_f64), pl.col("sec_dy_list").fill_null(empty_list_f64), pl.col("sec_dz_list").fill_null(empty_list_f64), ) .drop("_step_row") ) def _batch_to_polars(batch: ak.Array) -> pl.DataFrame: """Convert one awkward-array batch to a Polars DataFrame. Flat numeric/string fields are converted via numpy; variable-length fields (like child_track_ids) fall back to Python lists so polars stores them as List columns — a type parquet understands natively. """ col_dict: dict = {} for field in ak.fields(batch): arr = batch[field] if arr.ndim == 1 and not isinstance(arr.layout, ak.contents.ListOffsetArray): col_dict[field] = ak.to_numpy(arr) else: col_dict[field] = ak.to_list(arr) return pl.DataFrame(col_dict) def convert_steps_to_parquet( root_path: str | Path, output_path: str | Path | None = None, batch_size: str = "100 MB", tree_name: str = "Steps", compression: ParquetCompression = "snappy", ) -> Path: """Read *tree_name* from *root_path* and write it to a Parquet file. Reads in batches of *batch_size* so that peak ROOT-deserialization memory stays bounded. All batches are collected as Polars DataFrames and written in a single pass at the end (Polars' parquet writer does not support row-group appending without pyarrow). Parameters ---------- root_path: Input ROOT file. output_path: Output Parquet file. Defaults to *root_path* with .parquet suffix. batch_size: Uproot read batch size — an uproot size string ("100 MB") or integer row count (500_000). tree_name: Name of the TTree inside the ROOT file. compression: Parquet compression codec (snappy | lz4 | zstd | gzip | none). """ root_path = Path(root_path) if output_path is None: output_path = root_path.with_suffix(".parquet") else: output_path = Path(output_path) with uproot.open(root_path) as f: tree = f[tree_name] n_entries = tree.num_entries print(f"Reading '{tree_name}' from {root_path.name} ({n_entries} entries)") batches: list[pl.DataFrame] = [] rows_done = 0 for batch in tree.iterate(library="ak", step_size=batch_size): batches.append(_batch_to_polars(batch)) rows_done += len(batch) print(f" {rows_done:,} / {n_entries:,} rows read", end="\r", flush=True) df = pl.concat(batches) # Steps tree carries the parent→child links needed to derive secondary energy; # other trees (e.g. Hits) don't, so only augment when the column is present. if "child_track_ids" in df.columns: print("\nComputing per-step secondary attributes …", end=" ", flush=True) df = _add_secondary_attributes(df) print(f"\nWriting {output_path} …", end=" ", flush=True) df.write_parquet(output_path, compression=compression) print(f"done ({output_path.stat().st_size / 1e6:.1f} MB)") return output_path def main() -> None: parser = argparse.ArgumentParser( description="Convert a Steps (or any flat+jagged) tree in a ROOT file to Parquet." ) parser.add_argument("root_files", nargs="+", help="Input ROOT file(s)") parser.add_argument( "-o", "--output", help="Output Parquet file (default: .parquet). " "Only valid with a single input file.", ) 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)", ) args = parser.parse_args() if args.output is not None and len(args.root_files) > 1: parser.error("--output can only be used with a single input file") compression = "uncompressed" if args.compression == "none" else args.compression for root_file in args.root_files: convert_steps_to_parquet( root_file, output_path=args.output, batch_size=args.batch_size, tree_name=args.tree, compression=compression, ) if __name__ == "__main__": main()