Refactor VortexWrite to return Stream<Buffer> and Footer (#4557)

I attempted this in #4549 , but it was hanging, which has been resolved
by switching over to using the new runtimes API for execution. I could
have made this change separately, but both this PR and switching to
runtimes adds an argument to write_stream function so would have been a
conflict nightmare 😢

---

This revives an old branch I had worked on with @onursatici about
changing the Vortex write API to return a Stream and be completely
independent from the actual file sink.

This PR exposes a new write_stream function, and we can remove the old
write(VortexWrite, ArrayStream) APIs in a future PR. We should decide
whether we want to make that change prior to releasing to avoid multiple
user-facing breaks to this Rust API.

In order to make this change easier, we now pass an explicit eof-of-file
pointer into the writers that they can use to create new streams that
occur "just before the end of the file", as is required for zone maps
for example.

We also remove the custom transpose_stream implementation that was quite
dependent on the polling behaviour of `try_join_all`, which alarmingly
changes when the number of futures reaches 40.

Closes #4549 
Fixes #4552 
Fixes #4064

---------

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Author: gatesn

Cargo.lock

@@ -195,7 +195,7 @@ pub async fn convert_parquet_to_vortex(
                             compaction
                         );
                         let array_stream = parquet_to_vortex(parquet_file_path)?;
-                        let f = OpenOptions::new()
+                        let mut f = OpenOptions::new()
                             .write(true)
                             .truncate(true)
                             .create(true)
@@ -204,7 +204,7 @@ pub async fn convert_parquet_to_vortex(
 
                         let write_options = compaction.apply_options(VortexWriteOptions::default());
 
-                        write_options.write(f, array_stream).await?;
+                        write_options.write_tokio(&mut f, array_stream).await?;
 
                         anyhow::Ok(())
                     })

bench-vortex/src/clickbench/clickbench_data.rs

@@ -5,21 +5,16 @@ use std::io::Cursor;
 use std::sync::Arc;
 
 use bytes::Bytes;
-use tokio::runtime::Handle;
 use vortex::Array;
-use vortex::file::{VortexOpenOptions, VortexWriteOptions, WriteStrategyBuilder};
+use vortex::file::{VortexOpenOptions, VortexWriteOptions};
 
 #[inline(never)]
 pub async fn vortex_compress_write(array: &dyn Array, buf: &mut Vec<u8>) -> anyhow::Result<u64> {
-    Ok(VortexWriteOptions::default()
-        .with_strategy(
-            WriteStrategyBuilder::new()
-                .with_executor(Arc::new(Handle::current()))
-                .build(),
-        )
-        .write(Cursor::new(buf), array.to_array_stream())
-        .await?
-        .position())
+    let mut cursor = Cursor::new(buf);
+    VortexWriteOptions::default()
+        .write_tokio(&mut cursor, array.to_array_stream())
+        .await?;
+    Ok(cursor.position())
 }
 
 #[inline(never)]

bench-vortex/src/compress/vortex.rs

@@ -49,12 +49,14 @@ pub async fn fetch_taxi_data() -> Result<ArrayRef> {
 pub async fn taxi_data_vortex() -> Result<PathBuf> {
     idempotent_async("taxi.vortex", |output_fname| async move {
         let buf = output_fname.to_path_buf();
-        let output_file = File::create(output_fname).await?;
+        let mut output_file = File::create(output_fname).await?;
         VortexWriteOptions::default()
-            .write(output_file, parquet_to_vortex(taxi_data_parquet().await?)?)
-            .await?
-            .flush()
+            .write_tokio(
+                &mut output_file,
+                parquet_to_vortex(taxi_data_parquet().await?)?,
+            )
             .await?;
+        output_file.flush().await?;
         Ok(buf)
     })
     .await

bench-vortex/src/datasets/taxi_data.rs

@@ -27,12 +27,11 @@ use log::trace;
 use regex::Regex;
 use tokio::fs::File;
 use tokio::process::Command as TokioCommand;
-use tokio::runtime::Handle;
 use tracing::info;
 use url::Url;
 use vortex::ArrayRef;
 use vortex::error::{VortexResult, vortex_err};
-use vortex::file::{VortexOpenOptions, VortexWriteOptions, WriteStrategyBuilder};
+use vortex::file::{VortexOpenOptions, VortexWriteOptions};
 use vortex::iter::ArrayIteratorExt;
 use vortex::utils::aliases::hash_map::HashMap;
 use vortex_datafusion::VortexFormat;
@@ -346,13 +345,8 @@ impl PBIData {
                 let vortex_file =
                     idempotent_async(&vortex, async |output_path| -> anyhow::Result<()> {
                         VortexWriteOptions::default()
-                            .with_strategy(
-                                WriteStrategyBuilder::new()
-                                    .with_executor(Arc::new(Handle::current()))
-                                    .build(),
-                            )
-                            .write(
-                                File::create(output_path)
+                            .write_tokio(
+                                &mut File::create(output_path)
                                     .await
                                     .map_err(|e| anyhow::anyhow!("Failed to create file: {}", e))?,
                                 parquet_to_vortex(parquet)?,

bench-vortex/src/public_bi.rs

@@ -2,7 +2,6 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::io;
-use std::sync::Arc;
 
 use anyhow::{Context, Result, bail};
 use futures::StreamExt;
@@ -12,7 +11,6 @@ use parquet::arrow::{AsyncArrowWriter, ParquetRecordBatchStreamBuilder};
 use reqwest::Client;
 use tokio::fs::{File, create_dir_all};
 use tokio::io::BufReader;
-use tokio::runtime::Handle;
 use tokio_util::io::StreamReader;
 use tracing::info;
 use vortex::ArrayRef;
@@ -132,7 +130,7 @@ impl StatPopGenBenchmark {
 
     pub async fn parquet_to_vortex(&self, format: Format) -> Result<()> {
         let parquet_path = self.parquet_path()?;
-        let strategy = WriteStrategyBuilder::new().with_executor(Arc::new(Handle::current()));
+        let strategy = WriteStrategyBuilder::new();
         let (output_path, strategy) = match format {
             Format::OnDiskVortex => (self.vortex_path()?, strategy),
             Format::VortexCompact => (
@@ -179,8 +177,8 @@ impl StatPopGenBenchmark {
 
             VortexWriteOptions::default()
                 .with_strategy(strategy.build())
-                .write(
-                    File::create(output_path).await?,
+                .write_tokio(
+                    &mut File::create(output_path).await?,
                     ArrayStreamAdapter::new(dtype, vortex_stream),
                 )
                 .await?;

bench-vortex/src/statpopgen/download_vcf.rs

@@ -381,10 +381,10 @@ impl VortexWriter {
         let write_task = Some(tokio::spawn(async move {
             let stream = ArrayStreamAdapter::new(dtype, ReceiverStream::new(receiver));
 
-            let file = TokioFile::create(&file_path).await?;
+            let mut file = TokioFile::create(&file_path).await?;
             compaction_strategy
                 .apply_options(VortexWriteOptions::default())
-                .write(file, stream)
+                .write_tokio(&mut file, stream)
                 .await
                 .map_err(|e| anyhow!("Vortex write failed: {}", e))?;
 

bench-vortex/src/tpch/tpchgen.rs

@@ -20,6 +20,7 @@ vortex-array = { workspace = true }
 vortex-buffer = { workspace = true }
 vortex-dtype = { workspace = true }
 vortex-error = { workspace = true }
+vortex-io = { workspace = true }
 vortex-mask = { workspace = true }
 vortex-proto = { workspace = true }
 vortex-scalar = { workspace = true }

encodings/sequence/Cargo.toml

@@ -97,6 +97,7 @@ mod tests {
     use vortex_dtype::Nullability;
     use vortex_expr::{get_item, root};
     use vortex_file::{VortexOpenOptions, VortexWriteOptions};
+    use vortex_io::runtime::tokio::TokioRuntime;
     use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
 
     use crate::SequenceArray;
@@ -106,10 +107,10 @@ mod tests {
         let seq = SequenceArray::typed_new(2i8, 3, Nullability::NonNullable, 4).unwrap();
         let st = StructArray::from_fields(&[("a", seq.to_array())]).unwrap();
 
-        let file = tokio::fs::File::create("/tmp/abc.vx").await.unwrap();
+        let mut file = tokio::fs::File::create("/tmp/abc.vx").await.unwrap();
         VortexWriteOptions::default()
             .with_strategy(Arc::new(FlatLayoutStrategy::default()))
-            .write(file, st.to_array_stream())
+            .write(&mut file, st.to_array_stream(), TokioRuntime::handle())
             .await
             .unwrap();
 

encodings/sequence/src/serde.rs

@@ -31,6 +31,7 @@ vortex-dtype = { workspace = true, features = ["arbitrary"] }
 vortex-error = { workspace = true }
 vortex-expr = { workspace = true, features = ["arbitrary"] }
 vortex-file = { workspace = true, features = ["tokio"] }
+vortex-io = { workspace = true }
 vortex-mask = { workspace = true }
 vortex-scalar = { workspace = true, features = ["arbitrary"] }
 vortex-utils = { workspace = true }