polario.main

 1"""Main CLI interface
 2"""
 3import argparse
 4import json
 5import sys
 6from enum import Enum
 7from functools import reduce
 8from pathlib import Path
 9from pprint import pprint
10
11import polars as pl
12
13from polario import __version__
14
15
16class Command(Enum):
17    SHOW = "show"
18    SCHEMA = "schema"
19    JSON_HEAD = "json_head"
20    JSONL = "jsonl"
21    CONCAT_CSV = "concat_csv"
22    WRITE_CSV = "write_csv"
23
24
25def main() -> int:
26    """Main CLI interface"""
27    parser = argparse.ArgumentParser(
28        description="polario library commandline tool to inspect Parquet files"
29    )
30    parser.add_argument(
31        "--version",
32        action="version",
33        version="%(prog)s {version}".format(version=__version__),
34    )
35    parser.add_argument(
36        "cmd",
37        choices=[c.value for c in Command],
38        help="command to run",
39    )
40    parser.add_argument(
41        "paths",
42        metavar="PATH",
43        type=Path,
44        nargs="+",
45        help="input paths",
46    )
47    args = parser.parse_args()
48    cmd = Command(args.cmd)
49
50    if cmd == Command.CONCAT_CSV:
51        df = reduce(
52            lambda a, b: pl.concat([a, pl.read_csv(b, infer_schema_length=0)]),
53            args.paths[1:],
54            pl.read_csv(args.paths[0], infer_schema_length=0),
55        )
56        print(df)
57        output_filename = Path(args.paths[0].stem + ".parquet")
58        if output_filename.exists():
59            raise ValueError(f"Output file {output_filename} already exists")
60        print("Writing to", output_filename)
61        df.write_parquet(output_filename)
62        return 0
63    paths: list[Path] = args.paths
64    for path in paths:
65        if path.is_dir():
66            raise ValueError(
67                "Input path must be a file. File an issue if you want dataset support."
68            )
69        df = pl.read_parquet(path, use_pyarrow=True)
70        if cmd == Command.SHOW:
71            print(df)
72        elif cmd == Command.SCHEMA:
73            pprint(df.schema)
74        elif cmd == Command.JSON_HEAD:
75            json.dump(df.head().to_dicts(), sys.stdout, indent=2)
76        elif cmd == Command.JSONL:
77            for row in df.to_dicts():
78                json.dump(row, sys.stdout, separators=(",", ":"))
79                sys.stdout.write("\n")
80        elif cmd == Command.WRITE_CSV:
81            output_path = Path(path.name).with_suffix(".csv")
82            if output_path.exists():
83                print(f"Output file {output_path} already exists")
84                continue
85            df.write_csv(output_path)
86    return 0
87
88
89if __name__ == "__main__":
90    sys.exit(main())