Unify dataset/tooling scripts into a single dwarf Typer CLI
Replace the five separately-hyphenated uv entry points (steps-to-parquet, steps-to-parquet-parallel, migrate-geant-steps, bump-dataset-version, create-root-files) plus the unregistered hparam_scan.py with one `dwarf` command exposing convert/migrate/bump-gen/bump-schema/status/ update-manifest/create-manifest/make-root/hparam-scan as subcommands. Each scripts/*.py module now only holds argparse-free business logic; scripts/dwarf.py wires it up with Typer, matching giant/cli.py's style. `dwarf convert` merges the old serial/parallel conversion scripts behind a --jobs flag (default 1: sequential with plain -o; >1: dataset-layout fan-out via subprocess). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
@@ -1,14 +1,8 @@
|
||||
#!/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
|
||||
See `uv run dwarf convert --help` for the CLI.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
@@ -125,50 +119,3 @@ def convert_steps_to_parquet(
|
||||
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: <input>.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()
|
||||
|
||||
Reference in New Issue
Block a user