Allow steps_to_parquet.py to accept multiple ROOT input files

This commit is contained in:
2026-06-24 19:28:50 +02:00
parent eea9a24d9f
commit bdac36b203
+18 -9
View File
@@ -5,6 +5,7 @@ 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
@@ -86,9 +87,12 @@ 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("root_files", nargs="+", help="Input ROOT file(s)")
parser.add_argument(
"-o", "--output", help="Output Parquet file (default: <input>.parquet)"
"-o",
"--output",
help="Output Parquet file (default: <input>.parquet). "
"Only valid with a single input file.",
)
parser.add_argument(
"--batch-size",
@@ -108,13 +112,18 @@ def main() -> None:
)
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="uncompressed" if args.compression == "none" else args.compression,
)
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__":