feat: vortex CLI convert parquet to Vortex

Author: a10y

Cargo.lock

@@ -14,12 +14,17 @@ rust-version = { workspace = true }
 version = { workspace = true }
 
 [dependencies]
+arrow-array = { workspace = true }
 async-trait = { workspace = true }
 clap = { workspace = true, features = ["derive"] }
 crossterm = { workspace = true }
+futures-util = { workspace = true }
+indicatif = { workspace = true }
+parquet = { workspace = true, features = ["arrow", "async"] }
+pin-project = { workspace = true }
 ratatui = { workspace = true }
 tokio = { workspace = true, features = ["rt-multi-thread"] }
-vortex = { workspace = true, features = ["tokio"] }
+vortex = { workspace = true, features = ["tokio", "parquet"] }
 vortex-layout = { workspace = true }
 
 [lints]

vortex-tui/Cargo.toml

@@ -0,0 +1,88 @@
+use std::path::Path;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use arrow_array::StructArray as ArrowStructArray;
+use futures_util::Stream;
+use indicatif::ProgressBar;
+use parquet::arrow::ParquetRecordBatchStreamBuilder;
+use pin_project::pin_project;
+use tokio::fs::File;
+use vortex::arrays::ChunkedArray;
+use vortex::arrow::FromArrowArray;
+use vortex::dtype::DType;
+use vortex::error::{VortexExpect, VortexResult};
+use vortex::file::VortexWriteOptions;
+use vortex::stream::{ArrayStream, ArrayStreamArrayExt};
+use vortex::{Array, ArrayRef};
+
+/// Convert Parquet files to Vortex, and Vortex to Parquet.
+pub async fn exec_convert(input_path: impl AsRef<Path>) -> VortexResult<()> {
+    let output_path = input_path.as_ref().with_extension("vortex");
+    let file = File::open(input_path).await?;
+    let mut reader = ParquetRecordBatchStreamBuilder::new(file).await?.build()?;
+
+    // Create a new ChunkedArray and push a new chunk in for every RowGroup.
+
+    let mut chunks = Vec::new();
+
+    while let Some(mut reader) = reader.next_row_group().await? {
+        while let Some(batch) = reader.next() {
+            let batch = ArrowStructArray::from(batch?);
+            let next_chunk = ArrayRef::from_arrow(&batch, true);
+            chunks.push(next_chunk);
+        }
+    }
+
+    let dtype = chunks.first().vortex_expect("empty chunks").dtype().clone();
+    let chunked_array = ChunkedArray::try_new(chunks, dtype)?;
+
+    let writer = VortexWriteOptions::default();
+
+    let pb = ProgressBar::new(chunked_array.nchunks() as u64);
+    let stream = ProgressArrayStream {
+        inner: chunked_array.to_array_stream(),
+        progress: pb,
+    };
+
+    let output_file = File::create(output_path).await?;
+    writer.write(output_file, stream).await?;
+
+    Ok(())
+}
+
+#[pin_project]
+struct ProgressArrayStream<S> {
+    #[pin]
+    inner: S,
+    progress: ProgressBar,
+}
+
+impl<S> Stream for ProgressArrayStream<S>
+where
+    S: Stream<Item = VortexResult<ArrayRef>>,
+{
+    type Item = VortexResult<ArrayRef>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.project();
+        match this.inner.poll_next(cx) {
+            Poll::Ready(inner) => {
+                this.progress.inc(1);
+                Poll::Ready(inner)
+            }
+            Poll::Pending => {
+                return Poll::Pending;
+            }
+        }
+    }
+}
+
+impl<S> ArrayStream for ProgressArrayStream<S>
+where
+    S: ArrayStream,
+{
+    fn dtype(&self) -> &DType {
+        self.inner.dtype()
+    }
+}

vortex-tui/src/convert.rs

@@ -1,5 +1,6 @@
 #![allow(clippy::expect_used)]
 mod browse;
+mod convert;
 mod tree;
 
 use std::path::PathBuf;
@@ -10,6 +11,8 @@ use clap::Parser;
 use tokio::runtime::Runtime;
 use tree::exec_tree;
 
+use crate::convert::exec_convert;
+
 static TOKIO_RUNTIME: LazyLock<Runtime> =
     LazyLock::new(|| Runtime::new().expect("Tokio Runtime::new()"));
 
@@ -22,13 +25,17 @@ struct Cli {
 #[derive(Debug, clap::Subcommand)]
 enum Commands {
     Tree { file: PathBuf },
+    Convert { file: PathBuf },
     Browse { file: PathBuf },
 }
 
 fn main() {
     let cli = Cli::parse();
     match cli.command {
-        Commands::Tree { file } => exec_tree(file).expect("exec_tre"),
+        Commands::Tree { file } => TOKIO_RUNTIME.block_on(exec_tree(file)).expect("exec_tre"),
+        Commands::Convert { file } => TOKIO_RUNTIME
+            .block_on(exec_convert(file))
+            .expect("exec_convert"),
         Commands::Browse { file } => exec_tui(file).expect("exec_tui"),
     }
 }

vortex-tui/src/main.rs

@@ -4,19 +4,15 @@ use vortex::error::VortexResult;
 use vortex::file::VortexOpenOptions;
 use vortex::io::TokioFile;
 
-use crate::TOKIO_RUNTIME;
-
-pub fn exec_tree(file: impl AsRef<Path>) -> VortexResult<()> {
+pub async fn exec_tree(file: impl AsRef<Path>) -> VortexResult<()> {
     let opened = TokioFile::open(file)?;
 
-    let full = TOKIO_RUNTIME.block_on(async move {
-        VortexOpenOptions::file(opened)
-            .open()
-            .await?
-            .scan()
-            .into_array()
-            .await
-    })?;
+    let full = VortexOpenOptions::file(opened)
+        .open()
+        .await?
+        .scan()
+        .into_array()
+        .await?;
 
     println!("{}", full.tree_display());