Streaming conversion CLI (#2687)

Author: gatesn

Cargo.lock

@@ -14,14 +14,12 @@ 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 }
+indicatif = { workspace = true, features = ["futures"] }
 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", "parquet"] }

vortex-tui/Cargo.toml

@@ -1,30 +1,23 @@
-#![allow(clippy::use_debug)]
-
 use std::path::Path;
-use std::pin::Pin;
-use std::task::{Context, Poll};
-use std::time::Instant;
 
-use arrow_array::StructArray as ArrowStructArray;
-use futures_util::Stream;
+use futures_util::StreamExt;
 use indicatif::ProgressBar;
 use parquet::arrow::ParquetRecordBatchStreamBuilder;
-use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions};
-use pin_project::pin_project;
 use tokio::fs::File;
-use vortex::arrays::ChunkedArray;
-use vortex::arrow::FromArrowArray;
+use vortex::TryIntoArray;
 use vortex::dtype::DType;
-use vortex::error::{VortexExpect, VortexResult};
+use vortex::dtype::arrow::FromArrowType;
+use vortex::error::{VortexError, VortexExpect, VortexResult};
 use vortex::file::VortexWriteOptions;
-use vortex::stream::{ArrayStream, ArrayStreamArrayExt};
-use vortex::{Array, ArrayRef};
+use vortex::stream::ArrayStreamAdapter;
 
 #[derive(Default)]
 pub struct Flags {
     pub quiet: bool,
 }
 
+const BATCH_SIZE: usize = 8192;
+
 /// Convert Parquet files to Vortex.
 pub async fn exec_convert(input_path: impl AsRef<Path>, flags: Flags) -> VortexResult<()> {
     if !flags.quiet {
@@ -34,101 +27,41 @@ pub async fn exec_convert(input_path: impl AsRef<Path>, flags: Flags) -> VortexR
         );
     }
 
-    let wall_start = Instant::now();
-
     let output_path = input_path.as_ref().with_extension("vortex");
-    let mut file = File::open(input_path).await?;
-
-    let metadata =
-        ArrowReaderMetadata::load_async(&mut file, ArrowReaderOptions::default()).await?;
-    let has_root_level_nulls = metadata.parquet_schema().root_schema().is_optional();
-
-    let mut reader = ParquetRecordBatchStreamBuilder::new(file).await?.build()?;
-    let mut chunks = Vec::new();
-
-    while let Some(mut reader) = reader.next_row_group().await? {
-        for batch in reader.by_ref() {
-            let batch = ArrowStructArray::from(batch?);
-            let next_chunk = ArrayRef::from_arrow(&batch, has_root_level_nulls);
-            chunks.push(next_chunk);
-        }
-    }
-
-    let read_complete = Instant::now();
-
-    if !flags.quiet {
-        eprintln!(
-            "Read Parquet file in {:?}",
-            read_complete.duration_since(wall_start)
-        );
-
-        eprintln!(
-            "Writing {} chunks to new Vortex file {}",
-            chunks.len(),
-            output_path.display()
-        );
-    }
-
-    let dtype = chunks.first().vortex_expect("empty chunks").dtype().clone();
-    let chunked_array = ChunkedArray::try_new(chunks, dtype)?;
-
-    let writer = VortexWriteOptions::default();
-
-    let output_file = File::create(output_path).await?;
+    let file = File::open(input_path).await?;
+
+    let parquet = ParquetRecordBatchStreamBuilder::new(file)
+        .await?
+        .with_batch_size(BATCH_SIZE);
+    let num_rows = parquet.metadata().file_metadata().num_rows();
+
+    let dtype = DType::from_arrow(parquet.schema().as_ref());
+    let mut vortex_stream = parquet
+        .build()?
+        .map(|record_batch| {
+            record_batch
+                .map_err(VortexError::from)
+                .and_then(|rb| rb.try_into_array())
+        })
+        .boxed();
 
     if !flags.quiet {
-        let pb = ProgressBar::new(chunked_array.nchunks() as u64);
-        let stream = ProgressArrayStream {
-            inner: chunked_array.to_array_stream(),
-            progress: pb,
-        };
-        writer.write(output_file, stream).await?;
-    } else {
-        writer
-            .write(output_file, chunked_array.to_array_stream())
-            .await?;
+        // Parquet reader returns batches, rather than row groups. So make sure we correctly
+        // configure the progress bar.
+        let nbatches = u64::try_from(num_rows)
+            .vortex_expect("negative row count?")
+            .div_ceil(BATCH_SIZE as u64);
+        vortex_stream = ProgressBar::new(nbatches)
+            .wrap_stream(vortex_stream)
+            .boxed();
     }
 
-    if !flags.quiet {
-        eprintln!(
-            "Wrote Vortex in {:?}",
-            Instant::now().duration_since(read_complete)
-        );
-    }
+    VortexWriteOptions::default()
+        .write(
+            File::create(output_path).await?,
+            ArrayStreamAdapter::new(dtype, vortex_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 => Poll::Pending,
-        }
-    }
-}
-
-impl<S> ArrayStream for ProgressArrayStream<S>
-where
-    S: ArrayStream,
-{
-    fn dtype(&self) -> &DType {
-        self.inner.dtype()
-    }
-}