diff --git a/scripts/steps_to_parquet.py b/scripts/steps_to_parquet.py index 5d848f0..d7ccbb8 100644 --- a/scripts/steps_to_parquet.py +++ b/scripts/steps_to_parquet.py @@ -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: .parquet)" + "-o", + "--output", + help="Output Parquet file (default: .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__":