feat: jsonl output format, and multi-file input (#7)

Author: zachspar

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "parquet-py"
-version = "0.1.0-beta"
+version = "0.2.0-beta"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

README.md

@@ -9,7 +9,7 @@ Using Rust bindings under the hood, Parquet-Py provides a fast and efficient way
 ## Features
 
 - **Convert Parquet to CSV**: Convert your Parquet files into CSV format for easy viewing and processing in spreadsheet applications.
-- **Convert Parquet to JSON**: Easily convert your Parquet files into a JSON string format for quick inspection or processing.
+- **Convert Parquet to JSON / JSON Lines**: Easily convert your Parquet files into a JSON Array or JSON Lines format for quick inspection or processing.
 - **Iterable Parquet Rows**: Access Parquet file rows through an iterator, allowing for efficient row-by-row processing without loading the entire file into memory.
 - **Convert Parquet to Python List**: Transform your Parquet files into Python lists, where each row is represented as a dictionary within the list.
 
@@ -21,6 +21,16 @@ Using Rust bindings under the hood, Parquet-Py provides a fast and efficient way
 ## Usage
 ### Command-Line Interface
 
+> [!WARNING]
+> 
+> The CLI is still under development and may not be fully functional.
+> 
+> Breaking changes may occur in future releases.
+
+> [!TIP]
+> 
+> Multiple input files can be specified with `--input` option. For example, `--input file1.parquet --input file2.parquet`.
+
 #### Converting Parquet to CSV
 
 To convert a Parquet file into a CSV file, use the `parq convert` command.
@@ -29,14 +39,22 @@ To convert a Parquet file into a CSV file, use the `parq convert` command.
 parq convert --input path/to/your/file.parquet --format csv --output example.csv
 ```
 
-#### Converting Parquet to JSON
+#### Converting Parquet to JSON Array
 
-To convert a Parquet file into a JSON string, use the `parq convert` command.
+To convert a Parquet file into a JSON Array, use the `parq convert` command.
 
 ```bash
 parq convert --input path/to/your/file.parquet --format json --output example.json
 ```
 
+#### Converting Parquet to JSON Lines
+
+To convert a Parquet file into a JSON Lines, use the `parq convert` command.
+
+```bash
+parq convert --input path/to/your/file.parquet --format jsonl --output example.jsonl
+```
+
 
 ### Python
 

default.nix

@@ -6,14 +6,14 @@ let
       python3 -m venv env
       source env/bin/activate
       pip install -U pip
-      pip install 'maturin[patchelf]'
+      pip install 'maturin'
     else
-        source env/bin/activate
+      source env/bin/activate
     fi
   '';
 in
 
 pkgs.mkShell {
-    nativeBuildInputs = with pkgs.buildPackages; [ python312 ];
+    nativeBuildInputs = with pkgs.buildPackages; [ python312 cmake ninja autoconf automake rustc cargo rustfmt libiconv ];
     shellHook = createVenv;
 }

parq/cli/__main__.py

@@ -1,6 +1,11 @@
 #!/usr/bin/env python3
+from csv import DictWriter
+from json import dumps
+from io import StringIO
+from itertools import chain
+
 import click
-from parq import to_json_str, to_csv_str
+from parq import to_iter
 
 
 @click.group()
@@ -10,7 +15,13 @@ def parq_cli():
 
 @parq_cli.command()
 @click.option(
-    "--input", "-i", "parquet_file_path", help="Input parquet file", type=click.Path(exists=True), required=True
+    "--input",
+    "-i",
+    "parquet_file_paths",
+    help="Input parquet files",
+    type=click.Path(exists=True),
+    required=True,
+    multiple=True,
 )
 @click.option(
     "--format",
@@ -19,6 +30,7 @@ def parq_cli():
     type=click.Choice(
         [
             "json",
+            "jsonl",
             "csv",
         ],
         case_sensitive=False,
@@ -27,19 +39,55 @@ def parq_cli():
     help="Output format",
 )
 @click.option("--output", "-o", "output_file_path", help="Output file", type=click.Path(), required=False)
-def convert(parquet_file_path, output_format, output_file_path):
-    if output_format == "json":
-        output_string = to_json_str(str(parquet_file_path))
+def convert(parquet_file_paths, output_format, output_file_path):
+    """
+    Convert a list of parquet files to a specified output format.
+    :param parquet_file_paths: List of parquet file paths
+    :param output_format: Output format
+    :param output_file_path: Output file path
+    :return:
+    """
+
+    iterchain = chain.from_iterable([to_iter(parquet_file_path) for parquet_file_path in parquet_file_paths])
+
+    if output_format == "jsonl":
+        if output_file_path:
+
+            def _iter_jsonl(iterchain_):
+                for item_ in iterchain_:
+                    yield dumps(item_) + "\n"
+
+            with open(output_file_path, "w") as f:
+                f.writelines(_iter_jsonl(iterchain))
+        else:
+            for item in iterchain:
+                click.echo(dumps(item))
+
+    elif output_format == "json":
+        # FIXME: Optimize this
+        if output_file_path:
+            with open(output_file_path, "w") as f:
+                f.write(dumps(list(iterchain)))
+        else:
+            click.echo(dumps(list(iterchain)))
+
     elif output_format == "csv":
-        output_string = to_csv_str(str(parquet_file_path))
-    else:
-        raise ValueError(f"Unsupported output format {output_format}")
+        # FIXME: Optimize this
+        buffer = list(iterchain)
+        if output_file_path:
+            with open(output_file_path, "w") as f:
+                writer = DictWriter(f, fieldnames=buffer[0].keys())
+                writer.writeheader()
+                writer.writerows(buffer)
+        else:
+            f = StringIO()
+            writer = DictWriter(f, fieldnames=buffer[0].keys())
+            writer.writeheader()
+            writer.writerows(buffer)
+            click.echo(f.getvalue())
 
-    if output_file_path:
-        with open(output_file_path, "w") as f:
-            f.write(output_string)
     else:
-        click.echo(output_string)
+        raise ValueError(f"Unsupported output format: {output_format}")
 
 
 if __name__ == "__main__":