Allow steps_to_parquet.py to accept multiple ROOT input files
This commit is contained in:
@@ -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__":
|
||||
|
||||
Reference in New Issue
Block a user