Invert layout writers to push API

java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java

@@ -210,11 +210,10 @@ public synchronized long bytesWritten() {
     }
 
     /**
-     * Return the number of uncompressed bytes accepted by the writer but not yet written to the sink.
+     * Return the logical byte size of arrays currently retained by layout strategies.
      *
-     * <p>Together with {@link #bytesWritten()}, this lets callers estimate the in-progress file size: bytes that
-     * reached the sink are already compressed, while buffered bytes are still uncompressed and will shrink by roughly
-     * the file's observed compression ratio once flushed. After {@link #finish()}, this is zero.
+     * <p>This includes arrays queued for asynchronous layout work. It does not include allocator overhead,
+     * statistics-builder state, or buffering performed by the output sink. After {@link #finish()}, this is zero.
      */
     public synchronized long bufferedBytes() {
         if (summary != null) {

vortex-bench/src/tpch/tpchgen.rs

@@ -19,7 +19,6 @@ use parquet::file::properties::WriterProperties;
 use tokio::fs::File as TokioFile;
 use tokio::sync::Semaphore;
 use tokio::sync::mpsc;
-use tokio_stream::wrappers::ReceiverStream;
 use tokio_stream::wrappers::UnboundedReceiverStream;
 use tpchgen::generators::CustomerGenerator;
 use tpchgen::generators::LineItemGenerator;
@@ -31,9 +30,6 @@ use tpchgen::generators::RegionGenerator;
 use tpchgen::generators::SupplierGenerator;
 use tpchgen_arrow::RecordBatchIterator;
 use tracing::info;
-use vortex::array::ArrayRef;
-use vortex::array::stream::ArrayStreamAdapter;
-use vortex::error::VortexExpect;
 use vortex::file::WriteOptionsSessionExt;
 use vortex_arrow::ArrowSessionExt;
 
@@ -195,16 +191,12 @@ fn generate_table_file(
             // Create writer based on format
             let mut writer: Box<dyn FileWriter + Send> = match write_format {
                 Format::Parquet => Box::new(ParquetWriter::new(path, schema).await?),
-                Format::OnDiskVortex => Box::new(VortexWriter::new(
-                    path,
-                    schema,
-                    CompactionStrategy::Default,
-                )?),
-                Format::VortexCompact => Box::new(VortexWriter::new(
-                    path,
-                    schema,
-                    CompactionStrategy::Compact,
-                )?),
+                Format::OnDiskVortex => {
+                    Box::new(VortexWriter::new(path, schema, CompactionStrategy::Default).await?)
+                }
+                Format::VortexCompact => {
+                    Box::new(VortexWriter::new(path, schema, CompactionStrategy::Compact).await?)
+                }
                 _ => unreachable!(),
             };
 
@@ -324,37 +316,22 @@ impl FileWriter for ParquetWriter {
 
 /// Vortex writer for streaming TPC-H data
 struct VortexWriter {
-    sender: Option<mpsc::Sender<vortex::error::VortexResult<ArrayRef>>>,
-    write_task: Option<tokio::task::JoinHandle<Result<()>>>,
+    writer: vortex::file::Writer<TokioFile>,
 }
 
 impl VortexWriter {
-    fn new(
+    async fn new(
         path: PathBuf,
         schema: SchemaRef,
         compaction_strategy: CompactionStrategy,
     ) -> Result<Self> {
-        // Increase buffer size to avoid backpressure issues
-        let (sender, receiver) = mpsc::channel(2);
         let dtype = SESSION.arrow().from_arrow_schema(schema.as_ref())?;
-        let file_path = path;
-        let write_task = Some(tokio::spawn(async move {
-            let stream = ArrayStreamAdapter::new(dtype, ReceiverStream::new(receiver));
-
-            let mut file = TokioFile::create(&file_path).await?;
-            compaction_strategy
-                .apply_options(SESSION.write_options())
-                .write(&mut file, stream)
-                .await
-                .map_err(|e| anyhow!("Vortex write failed: {}", e))?;
-
-            Ok(())
-        }));
-
-        Ok(Self {
-            sender: Some(sender),
-            write_task,
-        })
+        let file = TokioFile::create(path).await?;
+        let writer = compaction_strategy
+            .apply_options(SESSION.write_options())
+            .writer(file, dtype)?;
+
+        Ok(Self { writer })
     }
 }
 
@@ -365,24 +342,17 @@ impl FileWriter for VortexWriter {
         let array = SESSION
             .arrow()
             .from_arrow_record_batch(batch.clone(), &schema)?;
-        self.sender
-            .as_ref()
-            .vortex_expect("sender closed early")
-            .send(Ok(array))
+        self.writer
+            .write(array)
             .await
-            .map_err(|_| anyhow!("Failed to send array to write task"))
+            .map_err(|e| anyhow!("Vortex write failed: {e}"))
     }
 
-    async fn finalize(mut self: Box<Self>) -> Result<()> {
-        // Close the sender to signal end of stream
-        self.sender.take();
-
-        // Wait for write task to complete
-        if let Some(task) = self.write_task.take() {
-            task.await
-                .map_err(|e| anyhow!("Write task failed: {}", e))??;
-        }
-
+    async fn finalize(self: Box<Self>) -> Result<()> {
+        self.writer
+            .close()
+            .await
+            .map_err(|e| anyhow!("Vortex write failed: {e}"))?;
         Ok(())
     }
 }

vortex-cuda/src/layout.rs

@@ -11,7 +11,6 @@ use std::sync::OnceLock;
 
 use async_trait::async_trait;
 use futures::FutureExt;
-use futures::StreamExt;
 use futures::future::BoxFuture;
 use vortex::array::ArrayRef;
 use vortex::array::ArrayVTable;
@@ -44,6 +43,7 @@ use vortex::layout::LayoutReader;
 use vortex::layout::LayoutReaderRef;
 use vortex::layout::LayoutRef;
 use vortex::layout::LayoutStrategy;
+use vortex::layout::LayoutWriter;
 use vortex::layout::LayoutWriterContext;
 use vortex::layout::RowSplits;
 use vortex::layout::SplitRange;
@@ -53,8 +53,7 @@ use vortex::layout::layouts::SharedArrayFuture;
 use vortex::layout::segments::SegmentId;
 use vortex::layout::segments::SegmentSinkRef;
 use vortex::layout::segments::SegmentSource;
-use vortex::layout::sequence::SendableSequentialStream;
-use vortex::layout::sequence::SequencePointer;
+use vortex::layout::sequence::SequenceId;
 use vortex::mask::Mask;
 use vortex::scalar::Scalar;
 use vortex::scalar::ScalarTruncation;
@@ -410,21 +409,40 @@ fn truncate_scalar_stat<F: Fn(Scalar) -> Option<(Scalar, bool)>>(
     }
 }
 
-#[async_trait]
 impl LayoutStrategy for CudaFlatLayoutStrategy {
-    async fn write_stream(
+    fn new_writer(
         &self,
         ctx: LayoutWriterContext,
         segment_sink: SegmentSinkRef,
-        mut stream: SendableSequentialStream,
-        _eof: SequencePointer,
+        dtype: DType,
         session: &VortexSession,
-    ) -> VortexResult<LayoutRef> {
-        let options = self.clone();
-        let Some(chunk) = stream.next().await else {
-            vortex_bail!("CudaFlatLayoutStrategy needs a single chunk");
-        };
-        let (sequence_id, chunk) = chunk?;
+    ) -> VortexResult<Box<dyn LayoutWriter>> {
+        Ok(Box::new(CudaFlatLayoutWriter {
+            ctx,
+            segment_sink,
+            dtype,
+            session: session.clone(),
+            options: self.clone(),
+            layout: None,
+        }))
+    }
+}
+
+struct CudaFlatLayoutWriter {
+    ctx: LayoutWriterContext,
+    segment_sink: SegmentSinkRef,
+    dtype: DType,
+    session: VortexSession,
+    options: CudaFlatLayoutStrategy,
+    layout: Option<LayoutRef>,
+}
+
+#[async_trait]
+impl LayoutWriter for CudaFlatLayoutWriter {
+    async fn write(&mut self, sequence_id: SequenceId, chunk: ArrayRef) -> VortexResult<()> {
+        if self.layout.is_some() {
+            vortex_bail!("CudaFlatLayoutStrategy received more than a single chunk");
+        }
         let row_count = chunk.len() as u64;
 
         match chunk.dtype() {
@@ -433,15 +451,15 @@ impl LayoutStrategy for CudaFlatLayoutStrategy {
                     lower_bound(
                         BufferString::from_scalar(v)
                             .vortex_expect("utf8 scalar must be a BufferString"),
-                        self.max_variable_length_statistics_size,
+                        self.options.max_variable_length_statistics_size,
                         *n,
                     )
                 });
                 truncate_scalar_stat(chunk.statistics(), Stat::Max, |v| {
                     upper_bound(
                         BufferString::from_scalar(v)
                             .vortex_expect("utf8 scalar must be a BufferString"),
-                        self.max_variable_length_statistics_size,
+                        self.options.max_variable_length_statistics_size,
                         *n,
                     )
                 });
@@ -451,15 +469,15 @@ impl LayoutStrategy for CudaFlatLayoutStrategy {
                     lower_bound(
                         ByteBuffer::from_scalar(v)
                             .vortex_expect("binary scalar must be a ByteBuffer"),
-                        self.max_variable_length_statistics_size,
+                        self.options.max_variable_length_statistics_size,
                         *n,
                     )
                 });
                 truncate_scalar_stat(chunk.statistics(), Stat::Max, |v| {
                     upper_bound(
                         ByteBuffer::from_scalar(v)
                             .vortex_expect("binary scalar must be a ByteBuffer"),
-                        self.max_variable_length_statistics_size,
+                        self.options.max_variable_length_statistics_size,
                         *n,
                     )
                 });
@@ -471,43 +489,52 @@ impl LayoutStrategy for CudaFlatLayoutStrategy {
         let host_buffers = extract_constant_buffers(&chunk);
 
         let buffers = chunk.serialize(
-            ctx.array_ctx(),
-            session,
+            self.ctx.array_ctx(),
+            &self.session,
             &SerializeOptions {
                 offset: 0,
-                include_padding: options.include_padding,
+                include_padding: self.options.include_padding,
             },
         )?;
         assert!(buffers.len() >= 2);
 
         // Always store the array tree inline (the cuda path requires it for planning).
         let array_tree = buffers[buffers.len() - 2].clone();
 
-        let segment_id = segment_sink.write(sequence_id, buffers).await?;
-
-        let None = stream.next().await else {
-            vortex_bail!("CudaFlatLayoutStrategy received stream with more than a single chunk");
-        };
+        let segment_id = self.segment_sink.write(sequence_id, buffers).await?;
 
         let host_buffer_map: HashMap<u32, ByteBuffer> = host_buffers
             .iter()
             .map(|hb| (hb.buffer_index, ByteBuffer::from(hb.data.clone())))
             .collect();
 
-        Ok(LayoutParts::new(
-            CudaFlat,
-            stream.dtype().clone(),
-            row_count,
-            vec![segment_id],
-            layout_children(Vec::new()),
-            CudaFlatData {
-                segment_id,
-                ctx: ReadContext::new(ctx.array_ctx().to_ids()),
-                array_tree,
-                host_buffers: Arc::new(host_buffer_map),
-            },
-        )
-        .into_layout())
+        self.layout = Some(
+            LayoutParts::new(
+                CudaFlat,
+                self.dtype.clone(),
+                row_count,
+                vec![segment_id],
+                layout_children(Vec::new()),
+                CudaFlatData {
+                    segment_id,
+                    ctx: ReadContext::new(self.ctx.array_ctx().to_ids()),
+                    array_tree,
+                    host_buffers: Arc::new(host_buffer_map),
+                },
+            )
+            .into_layout(),
+        );
+        Ok(())
+    }
+
+    async fn finish(&mut self, _sequence_id: SequenceId) -> VortexResult<()> {
+        Ok(())
+    }
+
+    async fn close(self: Box<Self>) -> VortexResult<LayoutRef> {
+        self.layout.ok_or_else(|| {
+            vortex::error::vortex_err!("CudaFlatLayoutStrategy needs a single chunk")
+        })
     }
 }
 

vortex-file/src/lib.rs

@@ -45,11 +45,14 @@
 //!
 //! # Writing
 //!
-//! Use [`WriteOptionsSessionExt::write_options`] or [`VortexWriteOptions::new`] to write an
-//! [`ArrayStream`](vortex_array::stream::ArrayStream). The default [`WriteStrategyBuilder`]
-//! repartitions rows, builds statistics layouts, dictionary-encodes suitable columns, compresses
-//! chunks with the BtrBlocks-style compressor, and writes flat leaf layouts. Advanced users can
-//! replace the whole strategy or override individual fields.
+//! Use [`WriteOptionsSessionExt::write_options`] or [`VortexWriteOptions::new`] to configure a
+//! write. For incremental writing, construct a [`Writer`] with [`VortexWriteOptions::writer`], call
+//! [`Writer::write`] for each array chunk, and finish with [`Writer::close`]. An
+//! [`ArrayStream`](vortex_array::stream::ArrayStream) can still be written in one operation with
+//! [`VortexWriteOptions::write`]. The default [`WriteStrategyBuilder`] repartitions rows, builds
+//! statistics layouts, dictionary-encodes suitable columns, compresses chunks with the
+//! BtrBlocks-style compressor, and writes flat leaf layouts. Advanced users can replace the whole
+//! strategy or override individual fields.
 //!
 //! # File Format
 //!

vortex-file/src/segments/writer.rs

@@ -54,7 +54,7 @@ impl SegmentSink for BufferedSegmentSink {
             let mut specs = self.segment_specs.lock();
             let segment_id = SegmentId::from(
                 u32::try_from(specs.len())
-                    .map_err(|_| vortex_err!("Too mant segments, u32 overflow"))?,
+                    .map_err(|_| vortex_err!("Too many segments, u32 overflow"))?,
             );
 
             // The API requires us to write these buffers contiguously. Therefore, we can only
@@ -90,10 +90,16 @@ impl SegmentSink for BufferedSegmentSink {
         };
 
         if let Some(padding) = padding_buffer {
-            let _ = self.buffers.send(padding).await;
+            self.buffers
+                .send(padding)
+                .await
+                .map_err(|_| vortex_err!("segment buffer receiver dropped"))?;
         }
         for buffer in buffers {
-            let _ = self.buffers.send(buffer).await;
+            self.buffers
+                .send(buffer)
+                .await
+                .map_err(|_| vortex_err!("segment buffer receiver dropped"))?;
         }
 
         Ok(segment_id)