Upgrade parquet, refactor, row count and negative offsets

Author: jupiter

Cargo.lock

@@ -1,7 +1,7 @@
 [package]
 name = "parquet2json"
-description = "A command-line tool for converting Parquet to newline-delimited JSON"
-version = "1.6.1"
+description = "A command-line tool for streaming Parquet as line-delimited JSON"
+version = "2.0.0"
 edition = "2018"
 license = "MIT"
 authors = ["Pieter Raubenheimer <[email protected]>"]
@@ -14,7 +14,7 @@ crossbeam-channel = { version = "0.5.1" }
 clap = "3.1.6"
 lazy_static = { version = "1.4.0" }
 openssl = { version = "0.10", features = ["vendored"] }
-parquet = { version = "10.0.0", features = ["cli"] }
+parquet = { version = "11.0.0", features = ["cli"] }
 regex = { version = "1.5.4" }
 reqwest = { version = "0.11.10", features = ["blocking"] }
 rusoto_core = { version = "0.47.0", default_features = false, features=["rustls"] }

Cargo.toml

@@ -1,8 +1,10 @@
 # parquet2json
 
-A command-line tool for converting [Parquet](https://parquet.apache.org) to [newline-delimited JSON](https://en.wikipedia.org/wiki/JSON_streaming#Line-delimited_JSON).
+A command-line tool for streaming [Parquet](https://parquet.apache.org) as [line-delimited JSON](https://en.wikipedia.org/wiki/JSON_streaming#Line-delimited_JSON).
 
-It uses the excellent [Apache Parquet Official Native Rust Implementation](https://github.com/apache/arrow-rs/tree/master/parquet).
+It reads only required ranges from file, HTTP or S3 locations, and supports offset/limit and column selection.
+
+It uses the [Apache Parquet Official Native Rust Implementation](https://github.com/apache/arrow-rs/tree/master/parquet) which has excellent support for compression formats and complex types.
 
 ## How to use it
 
@@ -13,18 +15,21 @@ $ cargo install parquet2json
 $ parquet2json --help
 
 USAGE:
-    parquet2json [OPTIONS] <FILE>
+    parquet2json [OPTIONS] <FILE> <SUBCOMMAND>
 
 ARGS:
     <FILE>    Location of Parquet input file (file path, HTTP or S3 URL)
 
 OPTIONS:
-    -o, --offset <OFFSET>                  Starts outputting from this row [default: 0]
-    -l, --limit <LIMIT>                    Maximum number of rows to output
-    -t, --timeout <TIMEOUT>                Request timeout in seconds [default: 60]
-    -s, --schema-output <SCHEMA_OUTPUT>    Outputs thrift schema only
-    -c, --columns <COLUMNS>                Select columns by name (comma,separated)
-    -h, --help                             Print help information
+    -t, --timeout <TIMEOUT>    Request timeout in seconds [default: 60]
+    -h, --help                 Print help information
+    -V, --version              Print version information
+
+SUBCOMMANDS:
+    cat         Outputs data as JSON lines
+    schema      Outputs the Thrift schema
+    rowcount    Outputs only the total row count
+    help        Print this message or the help of the given subcommand(s)
 ```
 
 ### S3 Settings
@@ -40,23 +45,23 @@ Use it to stream output to files and other tools such as `grep` and [jq](https:/
 #### Output to a file
 
 ```shell
-$ parquet2json ./myfile.pq > output.ndjson
+$ parquet2json ./myfile.pq cat > output.jsonl
 ```
 
-#### Filter with jq
+#### From S3 or HTTP (S3)
 
 ```shell
-$ parquet2json ./myfile.pq | jq 'select(.level==3) | .id'
+$ parquet2json s3://amazon-reviews-pds/parquet/product_category=Gift_Card/part-00000-495c48e6-96d6-4650-aa65-3c36a3516ddd.c000.snappy.parquet cat
 ```
 
-#### From S3 or HTTP (S3)
-
 ```shell
-$ parquet2json s3://amazon-reviews-pds/parquet/product_category=Gift_Card/part-00000-495c48e6-96d6-4650-aa65-3c36a3516ddd.c000.snappy.parquet
+$ parquet2json https://amazon-reviews-pds.s3.us-east-1.amazonaws.com/parquet/product_category%3DGift_Card/part-00000-495c48e6-96d6-4650-aa65-3c36a3516ddd.c000.snappy.parquet cat
 ```
 
+#### Filter selected columns with jq
+
 ```shell
-$ parquet2json https://amazon-reviews-pds.s3.us-east-1.amazonaws.com/parquet/product_category%3DGift_Card/part-00000-495c48e6-96d6-4650-aa65-3c36a3516ddd.c000.snappy.parquet
+$ parquet2json ./myfile.pq cat --columns=url,level | jq 'select(.level==3) | .url'
 ```
 
 ## License

README.md

@@ -1,9 +1,11 @@
 use core::time::Duration;
+use std::convert::TryInto;
 use std::fs::File;
+use std::ops::Add;
 use std::path::Path;
 use std::sync::Arc;
 
-use clap::{AppSettings, Parser};
+use clap::{AppSettings, Parser, Subcommand};
 use parquet::file::reader::{ChunkReader, FileReader, SerializedFileReader};
 use parquet::record::reader::RowIter;
 use parquet::schema::printer::print_schema;
@@ -101,25 +103,48 @@ fn get_projection<R: 'static + ChunkReader>(
     }
 }
 
-async fn print_json_from(
-    source: Source,
-    offset: u32,
-    limit: i32,
-    should_output_schema: bool,
-    timeout: Duration,
-    column_names: Option<String>,
+fn output_for_command<R: 'static + ChunkReader>(
+    file_reader: SerializedFileReader<R>,
+    command: &Commands,
 ) {
+    match command {
+        Commands::Cat {
+            offset,
+            limit,
+            columns,
+        } => {
+            let absolute_offset: u32 = if offset.is_negative() {
+                let parquet_metadata = file_reader.metadata();
+                parquet_metadata
+                    .file_metadata()
+                    .num_rows()
+                    .add(offset + 1)
+                    .try_into()
+                    .unwrap()
+            } else {
+                offset.abs().try_into().unwrap()
+            };
+            let projection = get_projection(&file_reader, columns.clone());
+            output_rows(&file_reader, projection, absolute_offset, *limit);
+        }
+        Commands::Schema {} => {
+            output_thrift_schema(&file_reader);
+        }
+        Commands::Rowcount {} => {
+            let parquet_metadata = file_reader.metadata();
+
+            println!("{}", parquet_metadata.file_metadata().num_rows());
+        }
+    }
+}
+
+async fn handle_command(source: Source, timeout: Duration, command: Commands) {
     match source {
         Source::File(path) => {
             let file = File::open(&Path::new(&path)).unwrap();
             let file_reader = SerializedFileReader::new(file).unwrap();
 
-            if should_output_schema {
-                output_thrift_schema(&file_reader);
-            } else {
-                let projection = get_projection(&file_reader, column_names);
-                output_rows(&file_reader, projection, offset, limit);
-            }
+            output_for_command(file_reader, &command);
         }
         Source::Http(url_str) => {
             let mut reader = HttpChunkReader::new_unknown_size(url_str).await;
@@ -128,12 +153,7 @@ async fn print_json_from(
             let blocking_task = tokio::task::spawn_blocking(move || {
                 let file_reader = SerializedFileReader::new(reader).unwrap();
 
-                if should_output_schema {
-                    output_thrift_schema(&file_reader);
-                } else {
-                    let projection = get_projection(&file_reader, column_names);
-                    output_rows(&file_reader, projection, offset, limit);
-                }
+                output_for_command(file_reader, &command);
             });
             blocking_task.await.unwrap();
         }
@@ -153,12 +173,7 @@ async fn print_json_from(
             let blocking_task = tokio::task::spawn_blocking(move || {
                 let file_reader = SerializedFileReader::new(reader).unwrap();
 
-                if should_output_schema {
-                    output_thrift_schema(&file_reader);
-                } else {
-                    let projection = get_projection(&file_reader, column_names);
-                    output_rows(&file_reader, projection, offset, limit);
-                }
+                output_for_command(file_reader, &command);
             });
             blocking_task.await.unwrap();
         }
@@ -172,67 +187,51 @@ struct Cli {
     /// Location of Parquet input file (file path, HTTP or S3 URL)
     file: String,
 
-    /// Starts outputting from this row
-    #[clap(default_value_t = 0, short, long, parse(try_from_str))]
-    offset: u32,
-
-    /// Maximum number of rows to output
-    #[clap(short, long, parse(try_from_str))]
-    limit: Option<i32>,
-
     /// Request timeout in seconds
     #[clap(default_value_t = 60, short, long, parse(try_from_str))]
     timeout: u16,
 
-    /// Outputs thrift schema only
-    #[clap(short, long)]
-    schema_output: Option<bool>,
+    #[clap(subcommand)]
+    command: Commands,
+}
+
+#[derive(Subcommand)]
+enum Commands {
+    /// Outputs data as JSON lines
+    Cat {
+        /// Starts outputting from this row (first row: 0, last row: -1)
+        #[clap(default_value_t = 0, short, long, parse(try_from_str))]
+        offset: i64,
+
+        /// Maximum number of rows to output
+        #[clap(short, long, parse(try_from_str), default_value_t = -1)]
+        limit: i32,
 
-    /// Select columns by name (comma,separated)
-    #[clap(short, long)]
-    columns: Option<String>,
+        /// Select columns by name (comma,separated)
+        #[clap(short, long)]
+        columns: Option<String>,
+    },
+
+    /// Outputs the Thrift schema
+    Schema {},
+
+    /// Outputs only the total row count
+    Rowcount {},
 }
 
 #[tokio::main]
 async fn main() {
     let cli = Cli::parse();
-
-    let output_thrift_schema = cli.schema_output.unwrap_or(false);
-    let offset = cli.offset;
-    let limit: i32 = cli.limit.unwrap_or(-1);
-    let timeout = Duration::from_secs(cli.timeout.into());
     let file = cli.file;
-    let column_names = cli.columns;
+    let timeout = Duration::from_secs(cli.timeout.into());
 
-    if file.as_str().starts_with("s3://") {
-        print_json_from(
-            Source::S3(file),
-            offset,
-            limit,
-            output_thrift_schema,
-            timeout,
-            column_names,
-        )
-        .await;
+    let source = if file.as_str().starts_with("s3://") {
+        Source::S3(file)
     } else if file.as_str().starts_with("http") {
-        print_json_from(
-            Source::Http(file),
-            offset,
-            limit,
-            output_thrift_schema,
-            timeout,
-            column_names,
-        )
-        .await;
+        Source::Http(file)
     } else {
-        print_json_from(
-            Source::File(file),
-            offset,
-            limit,
-            output_thrift_schema,
-            timeout,
-            column_names,
-        )
-        .await;
-    }
+        Source::File(file)
+    };
+
+    handle_command(source, timeout, cli.command).await;
 }