diff --git a/pyproject.toml b/pyproject.toml index 975b200..5faae65 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,11 @@ dependencies = [ dev = [ "pytest>=8", ] +convert = [ + "uproot>=5.3", + "awkward>=2.6", + "polars>=1.0", +] [project.scripts] giant = "giant.cli:app" diff --git a/scripts/steps_to_parquet.py b/scripts/steps_to_parquet.py new file mode 100644 index 0000000..61f3e14 --- /dev/null +++ b/scripts/steps_to_parquet.py @@ -0,0 +1,116 @@ +#!/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 +""" + +import argparse +from pathlib import Path + +import awkward as ak +import polars as pl +import uproot + + +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: str = "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) + + print(f"\nWriting {output_path} …", end=" ", flush=True) + pl.concat(batches).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_file", help="Input ROOT file") + parser.add_argument( + "-o", "--output", help="Output Parquet file (default: .parquet)" + ) + 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() + + convert_steps_to_parquet( + args.root_file, + output_path=args.output, + batch_size=args.batch_size, + tree_name=args.tree, + compression=None if args.compression == "none" else args.compression, + ) + + +if __name__ == "__main__": + main()