diff --git a/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java b/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java index c740a4b70f5..89faea0127a 100644 --- a/java/vortex-jni/src/main/java/dev/vortex/api/VortexWriter.java +++ b/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. * - *

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. + *

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) { diff --git a/vortex-bench/src/tpch/tpchgen.rs b/vortex-bench/src/tpch/tpchgen.rs index 2c4a26fd527..1c8c4c41918 100644 --- a/vortex-bench/src/tpch/tpchgen.rs +++ b/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 = 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>>, - write_task: Option>>, + writer: vortex::file::Writer, } impl VortexWriter { - fn new( + async fn new( path: PathBuf, schema: SchemaRef, compaction_strategy: CompactionStrategy, ) -> Result { - // 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) -> 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) -> Result<()> { + self.writer + .close() + .await + .map_err(|e| anyhow!("Vortex write failed: {e}"))?; Ok(()) } } diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 216b3369bd1..7e0de471eae 100644 --- a/vortex-cuda/src/layout.rs +++ b/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 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 { - let options = self.clone(); - let Some(chunk) = stream.next().await else { - vortex_bail!("CudaFlatLayoutStrategy needs a single chunk"); - }; - let (sequence_id, chunk) = chunk?; + ) -> VortexResult> { + 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, +} + +#[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,7 +451,7 @@ 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, ) }); @@ -441,7 +459,7 @@ impl LayoutStrategy for CudaFlatLayoutStrategy { 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,7 +469,7 @@ 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, ) }); @@ -459,7 +477,7 @@ impl LayoutStrategy for CudaFlatLayoutStrategy { 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,11 +489,11 @@ 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); @@ -483,31 +501,40 @@ impl LayoutStrategy for CudaFlatLayoutStrategy { // 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 = 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) -> VortexResult { + self.layout.ok_or_else(|| { + vortex::error::vortex_err!("CudaFlatLayoutStrategy needs a single chunk") + }) } } diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index 25760f0890e..a33c3542cda 100644 --- a/vortex-file/src/lib.rs +++ b/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 //! diff --git a/vortex-file/src/segments/writer.rs b/vortex-file/src/segments/writer.rs index e163c2cc868..7240ac5d3c7 100644 --- a/vortex-file/src/segments/writer.rs +++ b/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) diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index f5c177c9cdf..228265110e5 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -1476,7 +1476,7 @@ async fn test_writer_basic_push() -> VortexResult<()> { let dtype = st.dtype().clone(); let mut buf = ByteBufferMut::empty(); - let mut writer = SESSION.write_options().writer(&mut buf, dtype.clone()); + let mut writer = SESSION.write_options().writer(&mut buf, dtype.clone())?; writer.push(st.clone()).await?; let summary = writer.finish().await?; @@ -1492,6 +1492,29 @@ async fn test_writer_basic_push() -> VortexResult<()> { Ok(()) } +#[tokio::test] +async fn test_writer_rejects_mismatched_dtype() -> VortexResult<()> { + let array = buffer![1u32, 2, 3].into_array(); + let wrong_dtype = buffer![1i64, 2, 3].into_array(); + + let mut buf = ByteBufferMut::empty(); + let mut writer = SESSION + .write_options() + .writer(&mut buf, array.dtype().clone())?; + + let error = writer + .write(wrong_dtype) + .await + .expect_err("mismatched dtype must be rejected"); + assert!(error.to_string().contains("expected array with dtype u32")); + + writer.write(array).await?; + let summary = writer.close().await?; + assert_eq!(summary.row_count(), 3); + + Ok(()) +} + #[tokio::test] async fn test_writer_multiple_pushes() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); @@ -1505,7 +1528,7 @@ async fn test_writer_multiple_pushes() -> VortexResult<()> { let dtype = chunk1.dtype().clone(); let mut buf = ByteBufferMut::empty(); - let mut writer = SESSION.write_options().writer(&mut buf, dtype.clone()); + let mut writer = SESSION.write_options().writer(&mut buf, dtype.clone())?; writer.push(chunk1).await?; writer.push(chunk2).await?; @@ -1542,7 +1565,7 @@ async fn test_writer_push_stream() -> VortexResult<()> { let sendable_stream = ArrayStreamExt::boxed(ArrayStreamAdapter::new(dtype.clone(), stream)); let mut buf = ByteBufferMut::empty(); - let mut writer = SESSION.write_options().writer(&mut buf, dtype.clone()); + let mut writer = SESSION.write_options().writer(&mut buf, dtype.clone())?; writer.push_stream(sendable_stream).await?; @@ -1570,7 +1593,7 @@ async fn test_writer_bytes_written() -> VortexResult<()> { let dtype = array.dtype().clone(); let mut buf = ByteBufferMut::empty(); - let mut writer = SESSION.write_options().writer(&mut buf, dtype); + let mut writer = SESSION.write_options().writer(&mut buf, dtype)?; assert_eq!(writer.bytes_written(), 0); @@ -1592,16 +1615,15 @@ async fn test_writer_bytes_written() -> VortexResult<()> { } #[rstest] -#[case::table_one_leaf(true, 1, false, 32)] -#[case::table_two_shared_leaves(true, 2, false, 64)] -#[case::table_field_override(true, 2, true, 64)] -#[case::struct_default(false, 1, false, 32)] +#[case::table_one_leaf(true, 1, false)] +#[case::table_two_shared_leaves(true, 2, false)] +#[case::table_field_override(true, 2, true)] +#[case::struct_default(false, 1, false)] #[tokio::test] async fn test_writer_buffered_bytes( #[case] use_table_strategy: bool, #[case] leaf_count: usize, #[case] field_override: bool, - #[case] expected_buffered_bytes: u64, ) -> VortexResult<()> { const BUFFER_SIZE: u64 = 16; @@ -1631,17 +1653,23 @@ async fn test_writer_buffered_bytes( let mut buf = ByteBufferMut::empty(); let options = SESSION.write_options().with_strategy(Arc::clone(&strategy)); let buffered_bytes = options.buffered_bytes_tracker(); - let mut writer = options.writer(&mut buf, array.dtype().clone()); + let mut writer = options.writer(&mut buf, array.dtype().clone())?; assert_eq!(writer.buffered_bytes(), 0); - // The third push forces two chunks through the capacity-one input channel while keeping the - // writer open. Each physical leaf retains two BUFFER_SIZE chunks while peeking for more input. + // Each buffered leaf retains two chunks. Depending on how far its independently-driven child + // has progressed, the actor mailbox may also retain the third chunk. writer.push(array.clone()).await?; writer.push(array.clone()).await?; writer.push(array).await?; - assert_eq!(writer.buffered_bytes(), expected_buffered_bytes); + let observed_buffered_bytes = writer.buffered_bytes(); + let minimum = 2 * BUFFER_SIZE * leaf_count as u64; + let maximum = 3 * BUFFER_SIZE * leaf_count as u64; + assert!( + (minimum..=maximum).contains(&observed_buffered_bytes), + "expected {minimum}..={maximum} buffered bytes, got {observed_buffered_bytes}" + ); let summary = writer.finish().await?; assert_eq!(summary.row_count(), 12); @@ -1669,12 +1697,12 @@ async fn test_buffered_bytes_are_writer_scoped() -> VortexResult<()> { let mut first = SESSION .write_options() .with_strategy(Arc::clone(&strategy)) - .writer(&mut first_buf, array.dtype().clone()); + .writer(&mut first_buf, array.dtype().clone())?; let mut second_buf = ByteBufferMut::empty(); let mut second = SESSION .write_options() .with_strategy(strategy) - .writer(&mut second_buf, array.dtype().clone()); + .writer(&mut second_buf, array.dtype().clone())?; first.push(array.clone()).await?; first.push(array.clone()).await?; @@ -1683,8 +1711,8 @@ async fn test_buffered_bytes_are_writer_scoped() -> VortexResult<()> { second.push(array.clone()).await?; second.push(array).await?; - assert_eq!(first.buffered_bytes(), 2 * BUFFER_SIZE); - assert_eq!(second.buffered_bytes(), 2 * BUFFER_SIZE); + assert!((2 * BUFFER_SIZE..=3 * BUFFER_SIZE).contains(&first.buffered_bytes())); + assert!((2 * BUFFER_SIZE..=3 * BUFFER_SIZE).contains(&second.buffered_bytes())); first.finish().await?; second.finish().await?; @@ -1742,7 +1770,7 @@ async fn test_writer_empty_chunks() -> VortexResult<()> { let dtype = empty.dtype().clone(); let mut buf = ByteBufferMut::empty(); - let mut writer = SESSION.write_options().writer(&mut buf, dtype.clone()); + let mut writer = SESSION.write_options().writer(&mut buf, dtype.clone())?; writer.push(empty.clone()).await?; writer.push(non_empty).await?; @@ -1781,7 +1809,7 @@ async fn test_writer_mixed_push_and_stream() -> VortexResult<()> { let sendable_stream = ArrayStreamExt::boxed(ArrayStreamAdapter::new(dtype.clone(), stream)); let mut buf = ByteBufferMut::empty(); - let mut writer = SESSION.write_options().writer(&mut buf, dtype.clone()); + let mut writer = SESSION.write_options().writer(&mut buf, dtype.clone())?; writer.push(chunk1).await?; writer.push_stream(sendable_stream).await?; @@ -1824,7 +1852,7 @@ async fn test_writer_with_complex_types() -> VortexResult<()> { let dtype = chunk.dtype().clone(); let mut buf = ByteBufferMut::empty(); - let mut writer = SESSION.write_options().writer(&mut buf, dtype.clone()); + let mut writer = SESSION.write_options().writer(&mut buf, dtype.clone())?; writer.push(chunk).await?; let footer = writer.finish().await?; @@ -2020,7 +2048,7 @@ async fn test_writer_with_statistics() -> VortexResult<()> { let mut writer = SESSION .write_options() .with_file_statistics(PRUNING_STATS.to_vec()) - .writer(&mut buf, array.dtype().clone()); + .writer(&mut buf, array.dtype().clone())?; writer.push(array).await?; let summary = writer.finish().await?; diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 51261ce2628..7e81e6b0d2b 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -1,20 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::future::Future; use std::io; use std::io::Write; use std::sync::Arc; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; +use std::task::Poll; use futures::FutureExt; use futures::StreamExt; -use futures::TryStreamExt; -use futures::future::Fuse; -use futures::future::LocalBoxFuture; +use futures::future::poll_fn; use futures::future::ready; use futures::pin_mut; -use futures::select; use itertools::Itertools; use vortex_array::ArrayContext; use vortex_array::ArrayRef; @@ -25,30 +24,26 @@ use vortex_array::iter::ArrayIterator; use vortex_array::iter::ArrayIteratorExt; use vortex_array::stats::PRUNING_STATS; use vortex_array::stream::ArrayStream; -use vortex_array::stream::ArrayStreamAdapter; -use vortex_array::stream::ArrayStreamExt; use vortex_array::stream::SendableArrayStream; use vortex_buffer::ByteBuffer; use vortex_edition::ComponentKind; use vortex_edition::EditionSessionExt; -use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_io::IoBuf; use vortex_io::VortexWrite; -use vortex_io::kanal_ext::KanalExt; use vortex_io::runtime::BlockingRuntime; use vortex_io::session::RuntimeSessionExt; use vortex_layout::BufferedBytesTracker; use vortex_layout::LayoutContext; use vortex_layout::LayoutStrategy; +use vortex_layout::LayoutWriterActor; use vortex_layout::LayoutWriterContext; -use vortex_layout::layouts::file_stats::accumulate_stats; +use vortex_layout::layouts::file_stats::FileStatsAccumulator; use vortex_layout::sequence::SequenceId; -use vortex_layout::sequence::SequentialStreamAdapter; -use vortex_layout::sequence::SequentialStreamExt; +use vortex_layout::sequence::SequencePointer; use vortex_session::SessionExt; use vortex_session::VortexSession; use vortex_session::registry::Id; @@ -213,147 +208,67 @@ impl VortexWriteOptions { write: W, stream: S, ) -> VortexResult { - self.write_internal(write, ArrayStreamExt::boxed(stream)) - .await + let dtype = stream.dtype().clone(); + let mut writer = self.writer(write, dtype)?; + pin_mut!(stream); + while let Some(chunk) = stream.next().await { + writer.write(chunk?).await?; + } + writer.close().await } - async fn write_internal( - self, - mut write: W, - stream: SendableArrayStream, - ) -> VortexResult { + /// Create a push-based [`Writer`] that can be used to incrementally write arrays to the file. + /// + /// This follows the same lifecycle as other columnar file writers: call [`Writer::write`] for + /// each chunk, then call [`Writer::close`] to flush remaining buffers and receive the + /// [`WriteSummary`]. Each chunk must have dtype `dtype`. + pub fn writer(self, write: W, dtype: DType) -> VortexResult> { validate_metadata_segments(&self.metadata)?; - - // The array context is built here, rather than when the options were constructed, so that - // encodings registered on the session in between are still eligible for the file. let mut ctx = LayoutWriterContext::new(new_array_context(&self.session)) .with_buffered_bytes_tracker(self.buffered_bytes.clone()); if let Some(allowed) = edition_filter(&self.session, ComponentKind::Aggregate) { ctx = ctx.with_allowed_aggregates(allowed); } - let dtype = stream.dtype().clone(); - - let (mut ptr, eof) = SequenceId::root().split(); - - let stream = SequentialStreamAdapter::new( + let layout_ctx = new_layout_context(&self.session); + let (buffers_send, buffers) = kanal::bounded_async(1); + let segment_sink = Arc::new(BufferedSegmentSink::new( + buffers_send, + MAGIC_BYTES.len() as u64, + )); + let layout = self.strategy.new_writer( + ctx.clone(), + Arc::::clone(&segment_sink), dtype.clone(), - stream - .try_filter(|chunk| ready(!chunk.is_empty())) - .map(move |result| result.map(|chunk| (ptr.advance(), chunk))), - ) - .sendable(); - let (file_stats, stream) = accumulate_stats( - stream, + &self.session, + )?; + let layout = + LayoutWriterActor::spawn(layout, self.buffered_bytes.clone(), &self.session.handle()); + let sequence = SequenceId::root(); + let file_stats = FileStatsAccumulator::new( + &dtype, self.file_statistics.clone().into(), self.max_variable_length_statistics_size, &self.session, ); - - // First, write the magic bytes. - write.write_all(ByteBuffer::copy_from(MAGIC_BYTES)).await?; - let mut position = MAGIC_BYTES.len() as u64; - - // Create a channel to send buffers from the segment sink to the output stream. - let (send, recv) = kanal::bounded_async(1); - - let segments = Arc::new(BufferedSegmentSink::new(send, position)); - - // We spawn the layout future so it is driven in the background while we write the - // buffer stream, so we don't need to poll it until all buffers have been drained. - let ctx2 = ctx.clone(); - let session = self.session.clone(); - let layout_fut = self.session.handle().spawn_nested(move |h| async move { - let session = session.with_handle(h); - let layout = self - .strategy - .write_stream( - ctx2, - Arc::::clone(&segments), - stream, - eof, - &session, - ) - .await?; - Ok::<_, VortexError>((layout, segments.segment_specs())) - }); - - // Flush buffers as they arrive - let recv_stream = recv.into_stream(); - pin_mut!(recv_stream); - while let Some(buffer) = recv_stream.next().await { - if buffer.is_empty() { - continue; - } - position += buffer.len() as u64; - write.write_all(buffer).await?; - } - - let (layout, segment_specs) = layout_fut.await?; - - // Assemble the Footer object now that we have all the segments. - let statistics = if self.file_statistics.is_empty() { - None - } else { - Some(FileStatistics::new_with_dtype( - file_stats.stats_sets().into(), - &dtype, - )) - }; - let mut footer = Footer::new( - Arc::clone(&layout), - segment_specs, - statistics, - ReadContext::new(ctx.array_ctx().to_ids()), - ); - - // Emit the footer buffers and EOF. - let (footer_buffers, metadata, approx_byte_size) = footer - .clone() - .into_serializer() - .with_layout_context(new_layout_context(&self.session)) - .with_metadata_segments(self.metadata) - .with_offset(position) - .with_exclude_dtype(self.exclude_dtype) - .serialize_with_metadata()?; - footer = footer - .with_metadata_segments(metadata) - .with_approx_byte_size(approx_byte_size); - - for buffer in footer_buffers { - position += buffer.len() as u64; - write.write_all(buffer).await?; - } - - write.flush().await?; - - Ok(WriteSummary { - footer, - size: position, - }) - } - - /// Create a push-based [`Writer`] that can be used to incrementally write arrays to the file. - /// - /// Each pushed chunk must have dtype `dtype`. Call [`Writer::finish`] to close the input stream, - /// flush remaining buffers, and receive the [`WriteSummary`]. - pub fn writer<'w, W: VortexWrite + Unpin + 'w>(self, write: W, dtype: DType) -> Writer<'w> { - // Create a channel for sending arrays to the layout task. - let (arrays_send, arrays_recv) = kanal::bounded_async(1); - - let arrays = - ArrayStreamExt::boxed(ArrayStreamAdapter::new(dtype, arrays_recv.into_stream())); - let write = CountingVortexWrite::new(write); let bytes_written = write.counter(); - let buffered_bytes = self.buffered_bytes.clone(); - let future = self.write(write, arrays).boxed_local().fuse(); - - Writer { - arrays: Some(arrays_send), - future, + Ok(Writer { + write, + buffers, + segment_sink, + layout: Some(layout), + sequence, + ctx, + layout_ctx, + dtype, + file_stats, + file_statistics: self.file_statistics, + metadata: self.metadata, + exclude_dtype: self.exclude_dtype, + position: 0, bytes_written, - buffered_bytes, - } + buffered_bytes: self.buffered_bytes, + }) } } @@ -422,83 +337,153 @@ fn validate_metadata_segments(metadata: &HashMap) -> VortexR } /// An async API for writing Vortex files. -pub struct Writer<'w> { - // The input channel for sending arrays to the writer. - arrays: Option>>, - // The writer task that ultimately produces the footer. - future: Fuse>>, - // The bytes written so far. +pub struct Writer { + write: CountingVortexWrite, + buffers: kanal::AsyncReceiver, + segment_sink: Arc, + layout: Option, + sequence: SequencePointer, + ctx: LayoutWriterContext, + layout_ctx: LayoutContext, + dtype: DType, + file_stats: FileStatsAccumulator, + file_statistics: Vec, + metadata: HashMap, + exclude_dtype: bool, + position: u64, bytes_written: Arc, - // The buffered bytes accounting shared with the layout strategies for this write. buffered_bytes: BufferedBytesTracker, } -impl Writer<'_> { - /// Push a new chunk into the writer. - pub async fn push(&mut self, chunk: ArrayRef) -> VortexResult<()> { - let arrays = self.arrays.clone().vortex_expect("missing arrays sender"); - let send_fut = async move { arrays.send(Ok(chunk)).await }.fuse(); - pin_mut!(send_fut); - - // We poll the writer future to continue writing bytes to the output, while waiting for - // enough room to push the next chunk into the channel. - select! { - result = send_fut => { - // If the send future failed, the writer has failed or panicked. - if result.is_err() { - return Err(self.handle_failed_task().await); - } - }, - result = &mut self.future => { - // Under normal operation, the writer future should never complete until - // finish() is called. Therefore, we can assume the writer has failed. - // The writer future has failed, we need to propagate the error. - match result { - Ok(_) => vortex_bail!("Internal error: writer future completed early"), - Err(e) => return Err(e), - } - } +impl Writer { + async fn ensure_started(&mut self) -> VortexResult<()> { + if self.position == 0 { + self.write + .write_all(ByteBuffer::copy_from(MAGIC_BYTES)) + .await?; + self.position = MAGIC_BYTES.len() as u64; } + Ok(()) + } + async fn write_buffer( + write: &mut CountingVortexWrite, + position: &mut u64, + buffer: ByteBuffer, + ) -> VortexResult<()> { + if !buffer.is_empty() { + *position += buffer.len() as u64; + write.write_all(buffer).await?; + } Ok(()) } - /// Push an entire [`ArrayStream`] into the writer, consuming it. - /// - /// A task is spawned to consume the stream and push it into the writer, with the current - /// thread being used to write buffers to the output. - pub async fn push_stream(&mut self, mut stream: SendableArrayStream) -> VortexResult<()> { - let arrays = self.arrays.clone().vortex_expect("missing arrays sender"); - let stream_fut = async move { - while let Some(chunk) = stream.next().await { - arrays.send(chunk).await?; - } - Ok::<_, kanal::SendError>(()) + async fn drive_layout( + write: &mut CountingVortexWrite, + buffers: &kanal::AsyncReceiver, + position: &mut u64, + operation: impl Future>, + channel_closed_message: &'static str, + ) -> VortexResult { + enum Event { + Buffer(B), + Done, } - .fuse(); - pin_mut!(stream_fut); - - // We poll the writer future to continue writing bytes to the output, while waiting for - // enough room to push the stream into the channel. - select! { - result = stream_fut => { - if let Err(_send_err) = result { - // If the send future failed, the writer has failed or panicked. - return Err(self.handle_failed_task().await); + + let operation = operation.fuse(); + pin_mut!(operation); + let mut completed = None; + + loop { + let receive = buffers.recv().fuse(); + pin_mut!(receive); + let event = poll_fn(|cx| { + if let Poll::Ready(buffer) = receive.as_mut().poll(cx) { + return Poll::Ready(Event::Buffer(buffer)); + } + + if completed.is_none() + && let Poll::Ready(result) = operation.as_mut().poll(cx) + { + completed = Some(result); } - } - result = &mut self.future => { - // Under normal operation, the writer future should never complete until - // finish() is called. Therefore, we can assume the writer has failed. - // The writer future has failed, we need to propagate the error. - match result { - Ok(_) => vortex_bail!("Internal error: writer future completed early"), - Err(e) => return Err(e), + // Poll again because polling the layout operation may have completed a send into + // the receive future that was pending immediately above. + if let Poll::Ready(buffer) = receive.as_mut().poll(cx) { + return Poll::Ready(Event::Buffer(buffer)); + } + + if completed.is_some() { + Poll::Ready(Event::Done) + } else { + Poll::Pending + } + }) + .await; + + match event { + Event::Buffer(buffer) => { + let buffer = buffer.map_err(|_| vortex_err!("{channel_closed_message}"))?; + Self::write_buffer(write, position, buffer).await?; + } + Event::Done => { + return completed.take().vortex_expect("layout operation completed"); } } } + } + /// Write a new chunk. + /// + /// Returns an error without writing the chunk if its dtype does not match the dtype used to + /// construct the writer. + pub async fn write(&mut self, chunk: ArrayRef) -> VortexResult<()> { + if chunk.dtype() != &self.dtype { + vortex_bail!( + "Writer expected array with dtype {}, but received {}", + self.dtype, + chunk.dtype() + ); + } + + if chunk.is_empty() { + return Ok(()); + } + self.ensure_started().await?; + self.file_stats.push(&chunk)?; + + let sequence_id = self.sequence.advance(); + let layout = self.layout.as_mut().vortex_expect("layout writer present"); + let write = &mut self.write; + let buffers = &self.buffers; + let position = &mut self.position; + Self::drive_layout( + write, + buffers, + position, + layout.write(sequence_id, chunk), + "segment buffer channel closed while writing", + ) + .await?; + while let Ok(Some(buffer)) = buffers.try_recv() { + Self::write_buffer(write, position, buffer).await?; + } + Ok(()) + } + + /// Push a new chunk into the writer. + /// + /// This is an alias for [`Self::write`]. + pub async fn push(&mut self, chunk: ArrayRef) -> VortexResult<()> { + self.write(chunk).await + } + + /// Push an entire [`ArrayStream`] into the writer, consuming it. + pub async fn push_stream(&mut self, mut stream: SendableArrayStream) -> VortexResult<()> { + while let Some(chunk) = stream.next().await { + self.write(chunk?).await?; + } Ok(()) } @@ -507,7 +492,10 @@ impl Writer<'_> { self.bytes_written.load(Ordering::Relaxed) } - /// Returns the number of bytes currently buffered by the layout writers. + /// Returns the logical byte size of arrays currently retained by layout strategies. + /// + /// This includes arrays queued for asynchronous layout work. It does not include allocator + /// overhead, statistics-builder state, or buffering performed by the output sink. pub fn buffered_bytes(&self) -> u64 { self.buffered_bytes.buffered_bytes() } @@ -515,21 +503,67 @@ impl Writer<'_> { /// Finish writing the Vortex file, flushing any remaining buffers and returning the /// new file's footer. pub async fn finish(mut self) -> VortexResult { - // Drop the input channel to signal EOF. - drop(self.arrays.take()); + self.ensure_started().await?; + let mut layout = self.layout.take().vortex_expect("layout writer present"); + let sequence_id = self.sequence.advance(); + let write = &mut self.write; + let buffers = &self.buffers; + let position = &mut self.position; + let layout = Self::drive_layout( + write, + buffers, + position, + async move { + layout.finish(sequence_id).await?; + layout.take_layout() + }, + "segment buffer channel closed while closing", + ) + .await?; + while let Ok(Some(buffer)) = buffers.try_recv() { + Self::write_buffer(write, position, buffer).await?; + } - // Await the future task. - self.future.await + let statistics = if self.file_statistics.is_empty() { + None + } else { + Some(FileStatistics::new_with_dtype( + self.file_stats.stats_sets().into(), + &self.dtype, + )) + }; + let mut footer = Footer::new( + layout, + self.segment_sink.segment_specs(), + statistics, + ReadContext::new(self.ctx.array_ctx().to_ids()), + ); + let (footer_buffers, metadata, approx_byte_size) = footer + .clone() + .into_serializer() + .with_layout_context(self.layout_ctx) + .with_metadata_segments(self.metadata) + .with_offset(self.position) + .with_exclude_dtype(self.exclude_dtype) + .serialize_with_metadata()?; + footer = footer + .with_metadata_segments(metadata) + .with_approx_byte_size(approx_byte_size); + for buffer in footer_buffers { + Self::write_buffer(&mut self.write, &mut self.position, buffer).await?; + } + self.write.flush().await?; + Ok(WriteSummary { + footer, + size: self.position, + }) } - /// Assuming the writer task has failed, await it to get the error. - async fn handle_failed_task(&mut self) -> VortexError { - match (&mut self.future).await { - Ok(_) => vortex_err!( - "Internal error: writer task completed successfully but write future finished early" - ), - Err(e) => e, - } + /// Close the writer, flushing any remaining buffers and returning the file summary. + /// + /// This is an alias for [`Self::finish`]. + pub async fn close(self) -> VortexResult { + self.finish().await } } @@ -544,7 +578,7 @@ impl<'rt, B: BlockingRuntime> BlockingWrite<'rt, B> { /// /// The iterator is converted to an [`ArrayStream`] and driven to completion on /// the configured blocking runtime. - pub fn write( + pub fn write( self, write: W, iter: impl ArrayIterator + Send + 'static, @@ -557,25 +591,25 @@ impl<'rt, B: BlockingRuntime> BlockingWrite<'rt, B> { } /// Create a blocking push-based writer for chunks with dtype `dtype`. - pub fn writer<'w, W: Write + Unpin + 'w>( + pub fn writer( self, write: W, dtype: DType, - ) -> BlockingWriter<'rt, 'w, B> { - BlockingWriter { - writer: self.options.writer(BlockingWriteAdapter(write), dtype), + ) -> VortexResult> { + Ok(BlockingWriter { + writer: self.options.writer(BlockingWriteAdapter(write), dtype)?, runtime: self.runtime, - } + }) } } /// A blocking adapter around a [`Writer`], allowing incremental writing of arrays to a Vortex file. -pub struct BlockingWriter<'rt, 'w, B: BlockingRuntime> { +pub struct BlockingWriter<'rt, B: BlockingRuntime, W> { runtime: &'rt B, - writer: Writer<'w>, + writer: Writer>, } -impl BlockingWriter<'_, '_, B> { +impl BlockingWriter<'_, B, W> { /// Push one array chunk into the file. pub fn push(&mut self, chunk: ArrayRef) -> VortexResult<()> { self.runtime.block_on(self.writer.push(chunk)) @@ -586,7 +620,7 @@ impl BlockingWriter<'_, '_, B> { self.writer.bytes_written() } - /// Returns the number of bytes currently buffered by layout strategies. + /// Returns the logical byte size of arrays currently retained by layout strategies. pub fn buffered_bytes(&self) -> u64 { self.writer.buffered_bytes() } @@ -600,7 +634,7 @@ impl BlockingWriter<'_, '_, B> { // TODO(ngates): this blocking API may change, for now we just run blocking I/O inline. struct BlockingWriteAdapter(W); -impl VortexWrite for BlockingWriteAdapter { +impl VortexWrite for BlockingWriteAdapter { async fn write_all(&mut self, buffer: B) -> io::Result { self.0.write_all(buffer.as_slice())?; Ok(buffer) @@ -684,6 +718,12 @@ mod tests { use super::*; + #[test] + fn push_writer_is_send() { + fn assert_send() {} + assert_send::>>>(); + } + #[test] fn array_context_only_permits_enabled_encodings() -> Result<(), vortex_edition::EditionError> { const EDITION: EditionId = EditionId::new("test", 2026, 7, 0); diff --git a/vortex-io/src/write.rs b/vortex-io/src/write.rs index 9afc0a31015..3066bdf97ce 100644 --- a/vortex-io/src/write.rs +++ b/vortex-io/src/write.rs @@ -13,10 +13,10 @@ use vortex_buffer::ByteBufferMut; use crate::IoBuf; -pub trait VortexWrite { - fn write_all(&mut self, buffer: B) -> impl Future>; - fn flush(&mut self) -> impl Future>; - fn shutdown(&mut self) -> impl Future>; +pub trait VortexWrite: Send { + fn write_all(&mut self, buffer: B) -> impl Future> + Send; + fn flush(&mut self) -> impl Future> + Send; + fn shutdown(&mut self) -> impl Future> + Send; } impl VortexWrite for Vec { @@ -51,7 +51,7 @@ impl VortexWrite for ByteBufferMut { impl VortexWrite for Cursor where - Cursor: Write, + Cursor: Write + Send, { fn write_all(&mut self, buffer: B) -> impl Future> { ready(Write::write_all(self, buffer.as_slice()).map(|_| buffer)) @@ -112,7 +112,7 @@ impl VortexWrite for async_fs::File { /// An adapter to use an `AsyncWrite` as a `VortexWrite`. pub struct AsyncWriteAdapter(pub W); -impl VortexWrite for AsyncWriteAdapter { +impl VortexWrite for AsyncWriteAdapter { async fn write_all(&mut self, buffer: B) -> io::Result { self.0.write_all(buffer.as_slice()).await?; Ok(buffer) diff --git a/vortex-layout/src/layouts/buffered.rs b/vortex-layout/src/layouts/buffered.rs index b3a3e4652d9..e00fb86f920 100644 --- a/vortex-layout/src/layouts/buffered.rs +++ b/vortex-layout/src/layouts/buffered.rs @@ -4,21 +4,19 @@ use std::collections::VecDeque; use std::sync::Arc; -use async_stream::try_stream; use async_trait::async_trait; -use futures::StreamExt as _; -use futures::pin_mut; +use vortex_array::ArrayRef; +use vortex_array::dtype::DType; use vortex_error::VortexResult; use vortex_session::VortexSession; +use crate::BufferedBytesReservation; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriter; use crate::LayoutWriterContext; use crate::segments::SegmentSinkRef; -use crate::sequence::SendableSequentialStream; -use crate::sequence::SequencePointer; -use crate::sequence::SequentialStreamAdapter; -use crate::sequence::SequentialStreamExt as _; +use crate::sequence::SequenceId; #[derive(Clone)] pub struct BufferedStrategy { @@ -35,69 +33,83 @@ impl BufferedStrategy { } } -#[async_trait] impl LayoutStrategy for BufferedStrategy { - async fn write_stream( + fn new_writer( &self, ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, - stream: SendableSequentialStream, - eof: SequencePointer, + dtype: DType, session: &VortexSession, - ) -> VortexResult { - let dtype = stream.dtype().clone(); - let buffer_size = self.buffer_size; + ) -> VortexResult> { let buffered_bytes = ctx.buffered_bytes_tracker().clone(); + Ok(Box::new(BufferedLayoutWriter { + child: self.child.new_writer(ctx, segment_sink, dtype, session)?, + buffer_size: self.buffer_size, + buffered_bytes, + pending: None, + chunks: VecDeque::new(), + nbytes: 0, + })) + } +} - let buffered_stream = try_stream! { - let stream = stream.peekable(); - pin_mut!(stream); +struct BufferedLayoutWriter { + child: Box, + buffer_size: u64, + buffered_bytes: crate::BufferedBytesTracker, + pending: Option<(SequenceId, ArrayRef, BufferedBytesReservation)>, + chunks: VecDeque<(SequenceId, ArrayRef, BufferedBytesReservation)>, + nbytes: u64, +} - let mut nbytes = 0u64; - let mut chunks = VecDeque::new(); +impl BufferedLayoutWriter { + async fn process( + &mut self, + sequence_id: SequenceId, + chunk: ArrayRef, + reservation: BufferedBytesReservation, + last: bool, + ) -> VortexResult<()> { + self.nbytes += reservation.bytes(); + self.chunks.push_back((sequence_id, chunk, reservation)); - while let Some(chunk) = stream.as_mut().next().await { - let (sequence_id, chunk) = chunk?; - let chunk_size = chunk.nbytes(); - nbytes += chunk_size; - chunks.push_back((chunk, buffered_bytes.reserve(chunk_size))); + if !last && self.nbytes < 2 * self.buffer_size { + return Ok(()); + } - // If this is the last element, flush everything. - if stream.as_mut().peek().await.is_none() { - let mut sequence_ptr = sequence_id.descend(); - while let Some((chunk, reservation)) = chunks.pop_front() { - drop(reservation); - yield (sequence_ptr.advance(), chunk) - } - break; - } + while last || self.nbytes > self.buffer_size { + let Some((sequence_id, chunk, reservation)) = self.chunks.pop_front() else { + break; + }; + self.nbytes -= reservation.bytes(); + drop(reservation); + self.child.write(sequence_id, chunk).await?; + } + Ok(()) + } +} - if nbytes < 2 * buffer_size { - continue; - }; +#[async_trait] +impl LayoutWriter for BufferedLayoutWriter { + async fn write(&mut self, sequence_id: SequenceId, chunk: ArrayRef) -> VortexResult<()> { + let reservation = self.buffered_bytes.reserve(chunk.nbytes()); + if let Some((pending_id, pending, pending_reservation)) = + self.pending.replace((sequence_id, chunk, reservation)) + { + self.process(pending_id, pending, pending_reservation, false) + .await?; + } + Ok(()) + } - // Wait until we're at 2x the buffer size before flushing 1x the buffer size. - // This avoids small tail stragglers being flushed at the end of the file. - let mut sequence_ptr = sequence_id.descend(); - while nbytes > buffer_size { - let Some((chunk, reservation)) = chunks.pop_front() else { - break; - }; - nbytes -= reservation.bytes(); - drop(reservation); - yield (sequence_ptr.advance(), chunk) - } - } - }; + async fn finish(&mut self, sequence_id: SequenceId) -> VortexResult<()> { + if let Some((sequence_id, chunk, reservation)) = self.pending.take() { + self.process(sequence_id, chunk, reservation, true).await?; + } + self.child.finish(sequence_id).await + } - self.child - .write_stream( - ctx, - segment_sink, - SequentialStreamAdapter::new(dtype, buffered_stream).sendable(), - eof, - session, - ) - .await + async fn close(self: Box) -> VortexResult { + self.child.close().await } } diff --git a/vortex-layout/src/layouts/chunked/writer.rs b/vortex-layout/src/layouts/chunked/writer.rs index 8de7c30fe47..58b79c3a42b 100644 --- a/vortex-layout/src/layouts/chunked/writer.rs +++ b/vortex-layout/src/layouts/chunked/writer.rs @@ -3,26 +3,21 @@ use std::sync::Arc; -use async_stream::stream; use async_trait::async_trait; -use futures::StreamExt; -use futures::TryStreamExt; -use futures::stream; +use vortex_array::ArrayRef; +use vortex_array::dtype::DType; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriter; use crate::LayoutWriterContext; use crate::children::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; use crate::segments::SegmentSinkRef; -use crate::sequence::SendableSequentialStream; -use crate::sequence::SequencePointer; -use crate::sequence::SequentialStreamAdapter; -use crate::sequence::SequentialStreamExt as _; +use crate::sequence::SequenceId; #[derive(Clone)] pub struct ChunkedLayoutStrategy { @@ -38,62 +33,63 @@ impl ChunkedLayoutStrategy { } } -#[async_trait] impl LayoutStrategy for ChunkedLayoutStrategy { - async fn write_stream( + fn new_writer( &self, ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, - stream: SendableSequentialStream, - mut eof: SequencePointer, + dtype: DType, session: &VortexSession, - ) -> VortexResult { - let dtype = stream.dtype().clone(); - let dtype2 = dtype.clone(); - let chunk_strategy = Arc::clone(&self.chunk_strategy); - let handle = session.handle(); - - // We spawn each child to allow parallelism when processing chunks. - let stream = stream! { - let mut stream = stream; - while let Some(chunk) = stream.next().await { - let chunk_eof = eof.split_off(); + ) -> VortexResult> { + Ok(Box::new(ChunkedLayoutWriter { + chunk_strategy: Arc::clone(&self.chunk_strategy), + ctx, + segment_sink, + dtype, + session: session.clone(), + child_layouts: Vec::new(), + })) + } +} - let chunk_strategy = Arc::clone(&chunk_strategy); - let ctx = ctx.clone(); - let segment_sink = Arc::clone(&segment_sink); - let dtype = dtype2.clone(); - let session = session.clone(); +struct ChunkedLayoutWriter { + chunk_strategy: Arc, + ctx: LayoutWriterContext, + segment_sink: SegmentSinkRef, + dtype: DType, + session: VortexSession, + child_layouts: Vec, +} - yield handle.spawn_nested(move |handle| async move { - let session = session.with_handle(handle); - chunk_strategy - .write_stream( - ctx, - segment_sink, - SequentialStreamAdapter::new( - dtype, - stream::iter([chunk]), - ) - .sendable(), - chunk_eof, - &session, - ) - .await - }) - } - }; +#[async_trait] +impl LayoutWriter for ChunkedLayoutWriter { + async fn write(&mut self, sequence_id: SequenceId, chunk: ArrayRef) -> VortexResult<()> { + let mut sequence = sequence_id.descend(); + let mut child = self.chunk_strategy.new_writer( + self.ctx.clone(), + Arc::clone(&self.segment_sink), + self.dtype.clone(), + &self.session, + )?; + child.write(sequence.advance(), chunk).await?; + child.finish(sequence.advance()).await?; + self.child_layouts.push(child.close().await?); + Ok(()) + } - // Poll all of our children concurrently to accumulate their layouts. - let mut child_layouts: Vec = stream.buffered(usize::MAX).try_collect().await?; + async fn finish(&mut self, _sequence_id: SequenceId) -> VortexResult<()> { + Ok(()) + } + async fn close(mut self: Box) -> VortexResult { + let mut child_layouts = std::mem::take(&mut self.child_layouts); if child_layouts.len() == 1 { Ok(child_layouts.pop().vortex_expect("must have one child")) } else { let row_count = child_layouts.iter().map(|layout| layout.row_count()).sum(); Ok(ChunkedLayout::new( row_count, - dtype, + self.dtype, OwnedLayoutChildren::layout_children(child_layouts), ) .into_layout()) diff --git a/vortex-layout/src/layouts/collect.rs b/vortex-layout/src/layouts/collect.rs index 1fad03ef0a7..5bad747bee0 100644 --- a/vortex-layout/src/layouts/collect.rs +++ b/vortex-layout/src/layouts/collect.rs @@ -3,23 +3,22 @@ use std::sync::Arc; -use async_stream::try_stream; use async_trait::async_trait; -use futures::StreamExt; -use futures::pin_mut; +use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; +use vortex_array::dtype::DType; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; +use crate::BufferedBytesReservation; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriter; use crate::LayoutWriterContext; use crate::segments::SegmentSinkRef; -use crate::sequence::SendableSequentialStream; -use crate::sequence::SequencePointer; -use crate::sequence::SequentialStream; -use crate::sequence::SequentialStreamAdapter; +use crate::sequence::SequenceId; /// A strategy that collects all chunks and turns them into a single array chunk to pass into /// a child strategy. @@ -35,44 +34,58 @@ impl CollectStrategy { } } -#[async_trait] impl LayoutStrategy for CollectStrategy { - async fn write_stream( + fn new_writer( &self, ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, - stream: SendableSequentialStream, - eof: SequencePointer, + dtype: DType, session: &VortexSession, - ) -> VortexResult { - // Read the whole stream, then write one Chunked stream to the inner thing - let dtype = stream.dtype().clone(); - - let _dtype = dtype.clone(); - let collected_stream = try_stream! { - pin_mut!(stream); - - let mut chunks = Vec::new(); - let mut latest_sequence_id = None; - while let Some(chunk) = stream.next().await { - let (sequence_id, chunk) = chunk?; - latest_sequence_id = Some(sequence_id); - chunks.push(chunk); - } + ) -> VortexResult> { + let buffered_bytes = ctx.buffered_bytes_tracker().clone(); + Ok(Box::new(CollectLayoutWriter { + child: self + .child + .new_writer(ctx, segment_sink, dtype.clone(), session)?, + dtype, + buffered_bytes, + chunks: Vec::new(), + })) + } +} - // an empty input yields no chunk; the child layout handles it. - let Some(sequence_id) = latest_sequence_id else { - return; - }; +struct CollectLayoutWriter { + child: Box, + dtype: DType, + buffered_bytes: crate::BufferedBytesTracker, + chunks: Vec<(SequenceId, ArrayRef, BufferedBytesReservation)>, +} - let collected = ChunkedArray::try_new(chunks, _dtype)?.into_array(); - yield (sequence_id, collected); - }; +#[async_trait] +impl LayoutWriter for CollectLayoutWriter { + async fn write(&mut self, sequence_id: SequenceId, chunk: ArrayRef) -> VortexResult<()> { + let reservation = self.buffered_bytes.reserve(chunk.nbytes()); + self.chunks.push((sequence_id, chunk, reservation)); + Ok(()) + } - let adapted = Box::pin(SequentialStreamAdapter::new(dtype, collected_stream)); + async fn finish(&mut self, sequence_id: SequenceId) -> VortexResult<()> { + if !self.chunks.is_empty() { + let mut chunks = self.chunks.drain(..).collect::>(); + let (sequence_id, last, last_reservation) = + chunks.pop().vortex_expect("chunks checked non-empty"); + let chunks = chunks + .into_iter() + .map(|(_, chunk, _reservation)| chunk) + .chain(std::iter::once(last)); + drop(last_reservation); + let collected = ChunkedArray::try_new(chunks, self.dtype.clone())?.into_array(); + self.child.write(sequence_id, collected).await?; + } + self.child.finish(sequence_id).await + } - self.child - .write_stream(ctx, segment_sink, adapted, eof, session) - .await + async fn close(self: Box) -> VortexResult { + self.child.close().await } } diff --git a/vortex-layout/src/layouts/compressed.rs b/vortex-layout/src/layouts/compressed.rs index 87af23bb469..62c8fe249b9 100644 --- a/vortex-layout/src/layouts/compressed.rs +++ b/vortex-layout/src/layouts/compressed.rs @@ -1,13 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::collections::VecDeque; use std::sync::Arc; use async_trait::async_trait; -use futures::StreamExt as _; +use futures::FutureExt; +use futures::future::BoxFuture; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::VortexSessionExecute; +use vortex_array::dtype::DType; use vortex_array::expr::stats::Stat; use vortex_btrblocks::BtrBlocksCompressor; use vortex_error::VortexResult; @@ -15,14 +18,13 @@ use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; use vortex_utils::parallelism::get_available_parallelism; +use crate::BufferedBytesReservation; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriter; use crate::LayoutWriterContext; use crate::segments::SegmentSinkRef; -use crate::sequence::SendableSequentialStream; -use crate::sequence::SequencePointer; -use crate::sequence::SequentialStreamAdapter; -use crate::sequence::SequentialStreamExt; +use crate::sequence::SequenceId; /// A boxed compressor function from arrays into compressed arrays. /// @@ -85,46 +87,193 @@ impl CompressingStrategy { } } -#[async_trait] impl LayoutStrategy for CompressingStrategy { - async fn write_stream( + fn new_writer( &self, ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, - stream: SendableSequentialStream, - eof: SequencePointer, + dtype: DType, session: &VortexSession, - ) -> VortexResult { - let dtype = stream.dtype().clone(); + ) -> VortexResult> { + let buffered_bytes = ctx.buffered_bytes_tracker().clone(); + Ok(Box::new(CompressingLayoutWriter { + child: self.child.new_writer(ctx, segment_sink, dtype, session)?, + compressor: Arc::clone(&self.compressor), + stats: Arc::clone(&self.stats), + session: session.clone(), + buffered_bytes, + concurrency: self.concurrency, + pending: VecDeque::new(), + })) + } +} + +type CompressionFuture = + BoxFuture<'static, VortexResult<(SequenceId, ArrayRef, BufferedBytesReservation)>>; + +struct CompressingLayoutWriter { + child: Box, + compressor: Arc, + stats: Arc<[Stat]>, + session: VortexSession, + buffered_bytes: crate::BufferedBytesTracker, + concurrency: usize, + pending: VecDeque, +} + +impl CompressingLayoutWriter { + async fn drain_one(&mut self) -> VortexResult<()> { + let Some(result) = self.pending.pop_front() else { + return Ok(()); + }; + let (sequence_id, chunk, reservation) = result.await?; + drop(reservation); + self.child.write(sequence_id, chunk).await + } +} + +#[async_trait] +impl LayoutWriter for CompressingLayoutWriter { + async fn write(&mut self, sequence_id: SequenceId, chunk: ArrayRef) -> VortexResult<()> { let compressor = Arc::clone(&self.compressor); let stats = Arc::clone(&self.stats); - let session = session.clone(); - let compute_session = session.clone(); - - let handle = session.handle(); - let stream = stream - .map(move |chunk| { - let compressor = Arc::clone(&compressor); - let stats = Arc::clone(&stats); - let session = compute_session.clone(); - handle.spawn_cpu(move || { - let (sequence_id, chunk) = chunk?; + let session = self.session.clone(); + let reservation = self.buffered_bytes.reserve(chunk.nbytes()); + self.pending.push_back( + self.session + .handle() + .spawn_cpu(move || { let mut ctx = session.create_execution_ctx(); - // Compute the stats for the chunk prior to compression chunk.statistics().compute_all(&stats, &mut ctx)?; - Ok((sequence_id, compressor.compress_chunk(&chunk, &mut ctx)?)) + Ok(( + sequence_id, + compressor.compress_chunk(&chunk, &mut ctx)?, + reservation, + )) }) - }) - .buffered(self.concurrency); - - self.child - .write_stream( - ctx, - segment_sink, - SequentialStreamAdapter::new(dtype, stream).sendable(), - eof, - &session, - ) - .await + .boxed(), + ); + + if self.pending.len() >= self.concurrency { + self.drain_one().await?; + } + Ok(()) + } + + async fn finish(&mut self, sequence_id: SequenceId) -> VortexResult<()> { + while !self.pending.is_empty() { + self.drain_one().await?; + } + self.child.finish(sequence_id).await + } + + async fn close(self: Box) -> VortexResult { + self.child.close().await + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use std::time::Duration; + + use vortex_array::ArrayContext; + use vortex_array::IntoArray; + use vortex_array::arrays::StructArray; + use vortex_buffer::buffer; + use vortex_error::VortexExpect; + use vortex_io::session::RuntimeSessionExt; + + use super::*; + use crate::BufferedBytesTracker; + use crate::layouts::flat::writer::FlatLayoutStrategy; + use crate::layouts::struct_::StructStrategy; + use crate::segments::SegmentId; + use crate::segments::TestSegments; + use crate::test::SESSION; + use crate::test::new_session; + + #[tokio::test] + async fn spawned_compression_is_counted_as_buffered() -> VortexResult<()> { + let chunk = buffer![1u64, 2, 3, 4].into_array(); + let nbytes = chunk.nbytes(); + let dtype = chunk.dtype().clone(); + let tracker = BufferedBytesTracker::new(); + let ctx = LayoutWriterContext::new(ArrayContext::empty()) + .with_buffered_bytes_tracker(tracker.clone()); + let segments = Arc::new(TestSegments::default()); + let strategy = CompressingStrategy::new( + FlatLayoutStrategy::default(), + |chunk: &ArrayRef, _ctx: &mut ExecutionCtx| Ok(chunk.clone()), + ) + .with_concurrency(2) + .with_stats(&[]); + + let mut writer = strategy.new_writer(ctx, segments, dtype, &SESSION)?; + let mut sequence = SequenceId::root(); + writer.write(sequence.advance(), chunk).await?; + + assert_eq!(tracker.buffered_bytes(), nbytes); + + writer.finish(sequence.downgrade()).await?; + assert_eq!(tracker.buffered_bytes(), 0); + writer.close().await?; + Ok(()) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn struct_fields_compress_concurrently() -> VortexResult<()> { + let active = Arc::new(AtomicUsize::new(0)); + let max_active = Arc::new(AtomicUsize::new(0)); + let compressor = { + let active = Arc::clone(&active); + let max_active = Arc::clone(&max_active); + move |chunk: &ArrayRef, _ctx: &mut ExecutionCtx| { + let now = active.fetch_add(1, Ordering::SeqCst) + 1; + max_active.fetch_max(now, Ordering::SeqCst); + std::thread::sleep(Duration::from_millis(25)); + active.fetch_sub(1, Ordering::SeqCst); + Ok(chunk.clone()) + } + }; + let compressed: Arc = Arc::new( + CompressingStrategy::new(FlatLayoutStrategy::default(), compressor).with_stats(&[]), + ); + let strategy = StructStrategy::new(Arc::new(FlatLayoutStrategy::default()), compressed); + let chunk = StructArray::from_fields(&[ + ("a", buffer![1u64, 2, 3, 4].into_array()), + ("b", buffer![5u64, 6, 7, 8].into_array()), + ("c", buffer![9u64, 10, 11, 12].into_array()), + ("d", buffer![13u64, 14, 15, 16].into_array()), + ])? + .into_array(); + let session = new_session().with_tokio(); + let mut writer = strategy.new_writer( + LayoutWriterContext::new(ArrayContext::empty()), + Arc::new(TestSegments::default()), + chunk.dtype().clone(), + &session, + )?; + let mut sequence = SequenceId::root(); + + writer.write(sequence.advance(), chunk).await?; + writer.finish(sequence.downgrade()).await?; + let layout = writer.close().await?; + + assert!( + max_active.load(Ordering::SeqCst) > 1, + "compression did not overlap across struct fields" + ); + for (index, child) in layout.children()?.into_iter().enumerate() { + assert_eq!( + child.segment_ids(), + vec![SegmentId::from( + u32::try_from(index).vortex_expect("four fields fit in u32") + )] + ); + } + Ok(()) } } diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index b7c93a992fb..7fe67fba72e 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -1,23 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::pin::Pin; use std::sync::Arc; -use std::task::Context; -use std::task::Poll; -use async_stream::stream; -use async_stream::try_stream; use async_trait::async_trait; -use futures::FutureExt; -use futures::Stream; -use futures::StreamExt; -use futures::TryStreamExt; -use futures::future::BoxFuture; -use futures::pin_mut; -use futures::stream::BoxStream; -use futures::stream::once; -use futures::try_join; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -29,12 +15,9 @@ use vortex_array::builders::dict::dict_encoder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; -use vortex_io::kanal_ext::KanalExt; -use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; use crate::LayoutRef; @@ -45,12 +28,8 @@ use crate::layouts::chunked::ChunkedLayout; use crate::layouts::compressed::CompressorPlugin; use crate::layouts::dict::DictLayout; use crate::segments::SegmentSinkRef; -use crate::sequence::SendableSequentialStream; use crate::sequence::SequenceId; use crate::sequence::SequencePointer; -use crate::sequence::SequentialStream; -use crate::sequence::SequentialStreamAdapter; -use crate::sequence::SequentialStreamExt; /// Constraints for dictionary layout encoding. /// @@ -128,126 +107,226 @@ impl DictStrategy { } } -#[async_trait] impl LayoutStrategy for DictStrategy { - async fn write_stream( + fn new_writer( &self, ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, - stream: SendableSequentialStream, - mut eof: SequencePointer, + dtype: DType, session: &VortexSession, - ) -> VortexResult { - // Fallback if dtype is not supported - if !dict_layout_supported(stream.dtype()) { - return self - .fallback - .write_stream(ctx, segment_sink, stream, eof, session) - .await; + ) -> VortexResult> { + let mode = if dict_layout_supported(&dtype) { + None + } else { + Some(DictWriterMode::Fallback(self.fallback.new_writer( + ctx.clone(), + Arc::clone(&segment_sink), + dtype.clone(), + session, + )?)) + }; + Ok(Box::new(DictLayoutWriter { + codes: Arc::clone(&self.codes), + values: Arc::clone(&self.values), + fallback: Arc::clone(&self.fallback), + probe_compressor: Arc::clone(&self.probe_compressor), + constraints: self.options.constraints.clone().into(), + ctx, + segment_sink, + dtype, + session: session.clone(), + mode, + })) + } +} + +enum DictWriterMode { + Fallback(Box), + Dictionary(DictionaryLayoutWriter), +} + +struct DictLayoutWriter { + codes: Arc, + values: Arc, + fallback: Arc, + probe_compressor: Arc, + constraints: DictConstraints, + ctx: LayoutWriterContext, + segment_sink: SegmentSinkRef, + dtype: DType, + session: VortexSession, + mode: Option, +} + +impl DictLayoutWriter { + fn initialize(&mut self, first: &ArrayRef) -> VortexResult<()> { + let compressed = self + .probe_compressor + .compress_chunk(first, &mut self.session.create_execution_ctx())?; + self.mode = Some(if compressed.is::() { + DictWriterMode::Dictionary(DictionaryLayoutWriter { + codes: Arc::clone(&self.codes), + values: Arc::clone(&self.values), + ctx: self.ctx.clone(), + segment_sink: Arc::clone(&self.segment_sink), + dtype: self.dtype.clone(), + session: self.session.clone(), + encoder: DictStreamState { + encoder: None, + constraints: self.constraints.clone(), + }, + active_codes: None, + child_layouts: Vec::new(), + }) + } else { + DictWriterMode::Fallback(self.fallback.new_writer( + self.ctx.clone(), + Arc::clone(&self.segment_sink), + self.dtype.clone(), + &self.session, + )?) + }); + Ok(()) + } +} + +#[async_trait] +impl crate::LayoutWriter for DictLayoutWriter { + async fn write(&mut self, sequence_id: SequenceId, chunk: ArrayRef) -> VortexResult<()> { + if self.mode.is_none() { + self.initialize(&chunk)?; } + match self.mode.as_mut().vortex_expect("writer mode initialized") { + DictWriterMode::Fallback(writer) => writer.write(sequence_id, chunk).await, + DictWriterMode::Dictionary(writer) => writer.write(sequence_id, chunk).await, + } + } - let options = self.options.clone(); - let dtype = stream.dtype().clone(); - - // 0. decide if chunks are eligible for dict encoding - let (stream, first_chunk) = peek_first_chunk(stream).await?; - let stream = SequentialStreamAdapter::new(dtype.clone(), stream).sendable(); - - let should_fallback = match first_chunk { - None => true, // empty stream - Some(chunk) => { - let mut exec_ctx = session.create_execution_ctx(); - let compressed = self - .probe_compressor - .compress_chunk(&chunk, &mut exec_ctx)?; - !compressed.is::() - } - }; - if should_fallback { - // first chunk did not compress to dict, or did not exist. Skip dict layout - return self - .fallback - .write_stream(ctx, segment_sink, stream, eof, session) - .await; + async fn finish(&mut self, sequence_id: SequenceId) -> VortexResult<()> { + if self.mode.is_none() { + self.mode = Some(DictWriterMode::Fallback(self.fallback.new_writer( + self.ctx.clone(), + Arc::clone(&self.segment_sink), + self.dtype.clone(), + &self.session, + )?)); } + match self.mode.as_mut().vortex_expect("writer mode initialized") { + DictWriterMode::Fallback(writer) => writer.finish(sequence_id).await, + DictWriterMode::Dictionary(writer) => writer.finish(sequence_id).await, + } + } - // 1. from a chunk stream, create a stream that yields codes - // followed by a single value chunk when dict constraints are hit. - // (a1, a2) -> (code(c1), code(c2), values(v1), code(c3), ...) - let dict_stream = dict_encode_stream( - stream, - options.constraints.into(), - session.create_execution_ctx(), - ); + async fn close(mut self: Box) -> VortexResult { + match self.mode.take().vortex_expect("writer mode initialized") { + DictWriterMode::Fallback(writer) => writer.close().await, + DictWriterMode::Dictionary(writer) => Box::new(writer).close().await, + } + } +} + +struct DictionaryLayoutWriter { + codes: Arc, + values: Arc, + ctx: LayoutWriterContext, + segment_sink: SegmentSinkRef, + dtype: DType, + session: VortexSession, + encoder: DictStreamState, + active_codes: Option>, + child_layouts: Vec, +} - // Wrap up the dict stream to yield pairs of (codes_stream, values_future). - // Each of these pairs becomes a child dict layout. - let runs = DictionaryTransformer::new(dict_stream); - - let handle = session.handle(); - let dtype2 = dtype.clone(); - let child_layouts = stream! { - pin_mut!(runs); - - while let Some((codes_stream, values_fut)) = runs.next().await { - let codes = Arc::clone(&self.codes); - let codes_eof = eof.split_off(); - let ctx2 = ctx.clone(); - let segment_sink2 = Arc::clone(&segment_sink); - let session2 = session.clone(); - let codes_fut = handle.spawn_nested(move |h| async move { - let session2 = session2.with_handle(h); - codes.write_stream( - ctx2, - segment_sink2, - codes_stream.sendable(), - codes_eof, - &session2, - ).await - }); - - let values = Arc::clone(&self.values); - let values_eof = eof.split_off(); - let ctx2 = ctx.clone(); - let segment_sink2 = Arc::clone(&segment_sink); - let dtype2 = dtype2.clone(); - let session2 = session.clone(); - let values_layout = handle.spawn_nested(move |h| async move { - let session2 = session2.with_handle(h); - values.write_stream( - ctx2, - segment_sink2, - SequentialStreamAdapter::new(dtype2, once(values_fut)).sendable(), - values_eof, - &session2, - ).await - }); - - yield async move { - try_join!(codes_fut, values_layout) - }.boxed(); +impl DictionaryLayoutWriter { + async fn process(&mut self, chunk: DictionaryChunk) -> VortexResult<()> { + match chunk { + DictionaryChunk::Codes { + sequence_id, + codes, + codes_ptype, + } => { + if self.active_codes.is_none() { + self.active_codes = Some(self.codes.new_writer( + self.ctx.clone(), + Arc::clone(&self.segment_sink), + DType::Primitive(codes_ptype, Nullability::NonNullable), + &self.session, + )?); + } + self.active_codes + .as_mut() + .vortex_expect("codes writer active") + .write(sequence_id, codes) + .await } - }; + DictionaryChunk::Values(sequence_id, values) => { + let mut sequence = sequence_id.descend(); + let mut codes = self + .active_codes + .take() + .vortex_expect("values follow codes"); + let mut values_writer = self.values.new_writer( + self.ctx.clone(), + Arc::clone(&self.segment_sink), + self.dtype.clone(), + &self.session, + )?; + codes.finish(sequence.advance()).await?; + let codes_layout = codes.close().await?; + values_writer.write(sequence.advance(), values).await?; + values_writer.finish(sequence.advance()).await?; + let values_layout = values_writer.close().await?; + self.child_layouts + .push(DictLayout::new(values_layout, codes_layout).into_layout()); + Ok(()) + } + } + } - let mut child_layouts = child_layouts - .buffered(usize::MAX) - .map(|result| { - let (codes_layout, values_layout) = result?; - // All values are referenced when created via dictionary encoding - Ok::<_, VortexError>(DictLayout::new(values_layout, codes_layout).into_layout()) - }) - .try_collect::>() - .await?; + async fn write_chunk(&mut self, sequence_id: SequenceId, chunk: ArrayRef) -> VortexResult<()> { + let mut labeler = DictChunkLabeler::new(sequence_id); + let chunks = self.encoder.encode( + &mut labeler, + chunk, + &mut self.session.create_execution_ctx(), + )?; + for chunk in chunks { + self.process(chunk).await?; + } + Ok(()) + } +} - if child_layouts.len() == 1 { - return Ok(child_layouts.remove(0)); +#[async_trait] +impl crate::LayoutWriter for DictionaryLayoutWriter { + async fn write(&mut self, sequence_id: SequenceId, chunk: ArrayRef) -> VortexResult<()> { + self.write_chunk(sequence_id, chunk).await + } + + async fn finish(&mut self, sequence_id: SequenceId) -> VortexResult<()> { + let mut labeler = DictChunkLabeler::new(sequence_id); + for chunk in self.encoder.drain_values(&mut labeler) { + self.process(chunk).await?; } + if self.active_codes.is_some() { + return Err(vortex_err!("incomplete dictionary run")); + } + Ok(()) + } - let row_count = child_layouts.iter().map(|child| child.row_count()).sum(); + async fn close(mut self: Box) -> VortexResult { + if self.child_layouts.len() == 1 { + return Ok(self.child_layouts.pop().vortex_expect("one child layout")); + } + let row_count = self + .child_layouts + .iter() + .map(|child| child.row_count()) + .sum(); Ok(ChunkedLayout::new( row_count, - dtype, - OwnedLayoutChildren::layout_children(child_layouts), + self.dtype, + OwnedLayoutChildren::layout_children(self.child_layouts), ) .into_layout()) } @@ -255,57 +334,11 @@ impl LayoutStrategy for DictStrategy { enum DictionaryChunk { Codes { - seq_id: SequenceId, + sequence_id: SequenceId, codes: ArrayRef, codes_ptype: PType, }, - Values((SequenceId, ArrayRef)), -} - -type DictionaryStream = BoxStream<'static, VortexResult>; - -fn dict_encode_stream( - input: SendableSequentialStream, - constraints: DictConstraints, - mut exec_ctx: ExecutionCtx, -) -> DictionaryStream { - Box::pin(try_stream! { - let mut state = DictStreamState { - encoder: None, - constraints, - }; - - let input = input.peekable(); - pin_mut!(input); - - while let Some(item) = input.next().await { - let (sequence_id, chunk) = item?; - - // labeler potentially creates sub sequences, we must - // create it on both arms to avoid having a SequencePointer - // between await points - match input.as_mut().peek().await { - Some(_) => { - let mut labeler = DictChunkLabeler::new(sequence_id); - let chunks = state.encode(&mut labeler, chunk, &mut exec_ctx)?; - drop(labeler); - for dict_chunk in chunks { - yield dict_chunk; - } - } - None => { - // this is the last element, encode and drain chunks - let mut labeler = DictChunkLabeler::new(sequence_id); - let encoded = state.encode(&mut labeler, chunk, &mut exec_ctx)?; - let drained = state.drain_values(&mut labeler); - drop(labeler); - for dict_chunk in encoded.into_iter().chain(drained.into_iter()) { - yield dict_chunk; - } - } - } - } - }) + Values(SequenceId, ArrayRef), } struct DictStreamState { @@ -359,172 +392,34 @@ impl DictStreamState { } fn drain_values(&mut self, labeler: &mut DictChunkLabeler) -> Vec { - match self.encoder.as_mut() { + match self.encoder.take() { None => Vec::new(), - Some(encoder) => vec![labeler.values(encoder.reset())], + Some(mut encoder) => vec![labeler.values(encoder.reset())], } } } struct DictChunkLabeler { - sequence_pointer: SequencePointer, + sequence: SequencePointer, } impl DictChunkLabeler { - fn new(starting_id: SequenceId) -> Self { - let sequence_pointer = starting_id.descend(); - Self { sequence_pointer } - } - - fn codes(&mut self, chunk: ArrayRef, ptype: PType) -> DictionaryChunk { - DictionaryChunk::Codes { - seq_id: self.sequence_pointer.advance(), - codes: chunk, - codes_ptype: ptype, - } - } - - fn values(&mut self, chunk: ArrayRef) -> DictionaryChunk { - DictionaryChunk::Values((self.sequence_pointer.advance(), chunk)) - } -} - -type SequencedChunk = VortexResult<(SequenceId, ArrayRef)>; - -struct DictionaryTransformer { - input: DictionaryStream, - active_codes_tx: Option>, - active_values_tx: Option>, - pending_send: Option>>, -} - -impl DictionaryTransformer { - fn new(input: DictionaryStream) -> Self { + fn new(sequence_id: SequenceId) -> Self { Self { - input, - active_codes_tx: None, - active_values_tx: None, - pending_send: None, + sequence: sequence_id.descend(), } } -} -impl Stream for DictionaryTransformer { - type Item = (SendableSequentialStream, BoxFuture<'static, SequencedChunk>); - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - loop { - // First, try to complete any pending send - if let Some(mut send_fut) = self.pending_send.take() { - match send_fut.poll_unpin(cx) { - Poll::Ready(Ok(())) => { - // Send completed, continue processing - } - Poll::Ready(Err(_)) => { - // Receiver dropped, close this group - self.active_codes_tx = None; - if let Some(values_tx) = self.active_values_tx.take() { - drop(values_tx.send(Err(vortex_err!("values receiver dropped")))); - } - } - Poll::Pending => { - // Still pending, save it and return - self.pending_send = Some(send_fut); - return Poll::Pending; - } - } - } - - match self.input.poll_next_unpin(cx) { - Poll::Ready(Some(Ok(DictionaryChunk::Codes { - seq_id, - codes, - codes_ptype, - }))) => { - if self.active_codes_tx.is_none() { - // Start a new group - let (codes_tx, codes_rx) = kanal::bounded_async::(1); - let (values_tx, values_rx) = oneshot::channel(); - - self.active_codes_tx = Some(codes_tx.clone()); - self.active_values_tx = Some(values_tx); - - // Use passed codes_ptype instead of getting from array - let codes_dtype = DType::Primitive(codes_ptype, Nullability::NonNullable); - - // Send first codes. - self.pending_send = - Some(Box::pin( - async move { codes_tx.send(Ok((seq_id, codes))).await }, - )); - - // Create output streams. - let codes_stream = SequentialStreamAdapter::new( - codes_dtype, - codes_rx.into_stream().boxed(), - ) - .sendable(); - - let values_future = async move { - values_rx - .await - .map_err(|e| vortex_err!("values sender dropped: {}", e)) - .flatten() - } - .boxed(); - - return Poll::Ready(Some((codes_stream, values_future))); - } - - // Continue streaming codes to existing group - if let Some(tx) = &self.active_codes_tx { - let tx = tx.clone(); - self.pending_send = - Some(Box::pin(async move { tx.send(Ok((seq_id, codes))).await })); - } - } - Poll::Ready(Some(Ok(DictionaryChunk::Values(values)))) => { - // Complete the current group - if let Some(values_tx) = self.active_values_tx.take() { - drop(values_tx.send(Ok(values))); - } - self.active_codes_tx = None; // Close codes stream - } - Poll::Ready(Some(Err(e))) => { - // Send error to active channels if any - if let Some(values_tx) = self.active_values_tx.take() { - drop(values_tx.send(Err(e))); - } - self.active_codes_tx = None; - // And terminate the stream - return Poll::Ready(None); - } - Poll::Ready(None) => { - // Handle any incomplete group - if let Some(values_tx) = self.active_values_tx.take() { - drop(values_tx.send(Err(vortex_err!("Incomplete dictionary group")))); - } - self.active_codes_tx = None; - return Poll::Ready(None); - } - Poll::Pending => return Poll::Pending, - } + fn codes(&mut self, codes: ArrayRef, codes_ptype: PType) -> DictionaryChunk { + DictionaryChunk::Codes { + sequence_id: self.sequence.advance(), + codes, + codes_ptype, } } -} -async fn peek_first_chunk( - mut stream: BoxStream<'static, SequencedChunk>, -) -> VortexResult<(BoxStream<'static, SequencedChunk>, Option)> { - match stream.next().await { - None => Ok((stream.boxed(), None)), - Some(Err(e)) => Err(e), - Some(Ok((sequence_id, chunk))) => { - let chunk_clone = chunk.clone(); - let reconstructed_stream = - once(async move { Ok((sequence_id, chunk_clone)) }).chain(stream); - Ok((reconstructed_stream.boxed(), Some(chunk))) - } + fn values(&mut self, values: ArrayRef) -> DictionaryChunk { + DictionaryChunk::Values(self.sequence.advance(), values) } } @@ -589,33 +484,46 @@ fn remainder(array: &ArrayRef, encoded_len: usize) -> VortexResult = LazyLock::new(|| VortexSession::empty().with::()); - /// Regression test for a bug where the codes stream dtype was hardcoded to U16 instead of - /// using the actual codes dtype from the array. When `max_len <= 255`, the dict encoder - /// produces U8 codes, but the stream was incorrectly typed as U16, causing a dtype mismatch - /// assertion failure in [`SequentialStreamAdapter`]. - #[tokio::test] - async fn test_dict_transformer_uses_u8_for_small_dictionaries() { + fn encoded_codes_ptype( + arr: vortex_array::ArrayRef, + constraints: DictConstraints, + ) -> VortexResult { + let mut labeler = DictChunkLabeler::new(SequenceId::root().downgrade()); + let chunks = DictStreamState { + encoder: None, + constraints, + } + .encode(&mut labeler, arr, &mut SESSION.create_execution_ctx())?; + chunks + .into_iter() + .find_map(|chunk| match chunk { + DictionaryChunk::Codes { codes_ptype, .. } => Some(codes_ptype), + DictionaryChunk::Values(..) => None, + }) + .ok_or_else(|| vortex_err!("dictionary encoder produced no codes")) + } + + /// Regression test for selecting U8 codes when the configured dictionary fits in U8. + #[test] + fn test_dict_writer_uses_u8_for_small_dictionaries() -> VortexResult<()> { // Use max_len = 100 to force U8 codes (since 100 <= 255). let constraints = DictConstraints { max_bytes: 1024 * 1024, @@ -625,38 +533,17 @@ mod tests { // Create a simple string array with a few unique values. let arr = VarBinArray::from(vec!["hello", "world", "hello", "world"]).into_array(); - // Wrap into a sequential stream. - let mut pointer = SequenceId::root(); - let input_stream = SequentialStreamAdapter::new( - arr.dtype().clone(), - futures::stream::once(async move { Ok((pointer.advance(), arr)) }), - ) - .sendable(); - - // Encode into dict chunks. - let dict_stream = - dict_encode_stream(input_stream, constraints, SESSION.create_execution_ctx()); - - // Transform into codes/values streams. - let mut transformer = DictionaryTransformer::new(dict_stream); - - // Get the first (and only) run. - let (codes_stream, _values_fut) = transformer - .next() - .await - .expect("expected at least one dictionary run"); - - // The key assertion: codes stream dtype should be U8, not U16. assert_eq!( - codes_stream.dtype(), - &DType::Primitive(PType::U8, NonNullable), - "codes stream should use U8 dtype for small dictionaries, not U16" + encoded_codes_ptype(arr, constraints)?, + PType::U8, + "codes should use U8 for small dictionaries" ); + Ok(()) } - /// Test that the codes stream uses U16 dtype when the dictionary has more than 255 entries. - #[tokio::test] - async fn test_dict_transformer_uses_u16_for_large_dictionaries() { + /// Test that the codes use U16 when the dictionary may contain more than 255 entries. + #[test] + fn test_dict_writer_uses_u16_for_large_dictionaries() -> VortexResult<()> { // Use max_len = 1000 to allow U16 codes (since 1000 > 255). let constraints = DictConstraints { max_bytes: 1024 * 1024, @@ -668,32 +555,11 @@ mod tests { let arr = VarBinArray::from(values.iter().map(|s| s.as_str()).collect::>()).into_array(); - // Wrap into a sequential stream. - let mut pointer = SequenceId::root(); - let input_stream = SequentialStreamAdapter::new( - arr.dtype().clone(), - futures::stream::once(async move { Ok((pointer.advance(), arr)) }), - ) - .sendable(); - - // Encode into dict chunks. - let dict_stream = - dict_encode_stream(input_stream, constraints, SESSION.create_execution_ctx()); - - // Transform into codes/values streams. - let mut transformer = DictionaryTransformer::new(dict_stream); - - // Get the first (and only) run. - let (codes_stream, _values_fut) = transformer - .next() - .await - .expect("expected at least one dictionary run"); - - // Codes stream dtype should be U16 since we have more than 255 distinct values. assert_eq!( - codes_stream.dtype(), - &DType::Primitive(PType::U16, NonNullable), - "codes stream should use U16 dtype for dictionaries with >255 entries" + encoded_codes_ptype(arr, constraints)?, + PType::U16, + "codes should use U16 for large dictionaries" ); + Ok(()) } } diff --git a/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index dd0c4d6c13c..381a06f5f25 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -428,7 +428,7 @@ pub struct FileStatsAccumulator { } impl FileStatsAccumulator { - fn new( + pub fn new( dtype: &DType, stats: Arc<[Stat]>, max_variable_length_statistics_size: usize, @@ -475,6 +475,12 @@ impl FileStatsAccumulator { chunk: VortexResult<(SequenceId, ArrayRef)>, ) -> VortexResult<(SequenceId, ArrayRef)> { let (sequence_id, chunk) = chunk?; + self.push(&chunk)?; + Ok((sequence_id, chunk)) + } + + /// Accumulate statistics for one pushed array chunk. + pub fn push(&self, chunk: &ArrayRef) -> VortexResult<()> { let mut ctx = self.ctx.lock(); if chunk.dtype().is_struct() { let struct_chunk = chunk.clone().execute::(&mut ctx)?; @@ -487,9 +493,9 @@ impl FileStatsAccumulator { acc.push_chunk(field, &mut ctx)?; } } else { - self.accumulators.lock()[0].push_chunk(&chunk, &mut ctx)?; + self.accumulators.lock()[0].push_chunk(chunk, &mut ctx)?; } - Ok((sequence_id, chunk)) + Ok(()) } pub fn stats_sets(&self) -> Vec { diff --git a/vortex-layout/src/layouts/flat/writer.rs b/vortex-layout/src/layouts/flat/writer.rs index 9761c71f9ae..e0a1a29593a 100644 --- a/vortex-layout/src/layouts/flat/writer.rs +++ b/vortex-layout/src/layouts/flat/writer.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use async_trait::async_trait; -use futures::StreamExt; +use vortex_array::ArrayRef; use vortex_array::dtype::DType; use vortex_array::expr::stats::Precision; use vortex_array::expr::stats::Stat; @@ -23,14 +23,14 @@ use vortex_session::registry::ReadContext; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriter; use crate::LayoutWriterContext; use crate::children::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; use crate::layouts::flat::FlatLayout; use crate::layouts::flat::flat_layout_inline_array_node; use crate::segments::SegmentSinkRef; -use crate::sequence::SendableSequentialStream; -use crate::sequence::SequencePointer; +use crate::sequence::SequenceId; #[derive(Clone)] pub struct FlatLayoutStrategy { @@ -79,27 +79,42 @@ fn truncate_scalar_stat Option<(Scalar, bool)>>( } } -#[async_trait] impl LayoutStrategy for FlatLayoutStrategy { - async fn write_stream( + fn new_writer( &self, ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, - mut stream: SendableSequentialStream, - _eof: SequencePointer, + dtype: DType, session: &VortexSession, - ) -> VortexResult { - let Some(chunk) = stream.next().await else { - // an empty input has no segment to write. - return Ok(ChunkedLayout::new( - 0, - stream.dtype().clone(), - OwnedLayoutChildren::layout_children(vec![]), - ) - .into_layout()); - }; - let (sequence_id, chunk) = chunk?; + ) -> VortexResult> { + Ok(Box::new(FlatLayoutWriter { + ctx, + segment_sink, + dtype, + session: session.clone(), + include_padding: self.include_padding, + max_variable_length_statistics_size: self.max_variable_length_statistics_size, + layout: None, + })) + } +} +struct FlatLayoutWriter { + ctx: LayoutWriterContext, + segment_sink: SegmentSinkRef, + dtype: DType, + session: VortexSession, + include_padding: bool, + max_variable_length_statistics_size: usize, + layout: Option, +} + +#[async_trait] +impl LayoutWriter for FlatLayoutWriter { + async fn write(&mut self, sequence_id: SequenceId, chunk: ArrayRef) -> VortexResult<()> { + if self.layout.is_some() { + vortex_bail!("flat layout received more than a single chunk"); + } let row_count = chunk.len() as u64; match chunk.dtype() { @@ -143,8 +158,8 @@ impl LayoutStrategy for FlatLayoutStrategy { } let buffers = chunk.serialize( - ctx.array_ctx(), - session, + self.ctx.array_ctx(), + &self.session, &SerializeOptions { offset: 0, include_padding: self.include_padding, @@ -154,19 +169,29 @@ impl LayoutStrategy for FlatLayoutStrategy { assert!(buffers.len() >= 2); let array_node = flat_layout_inline_array_node().then(|| buffers[buffers.len() - 2].clone()); - let segment_id = segment_sink.write(sequence_id, buffers).await?; - - let None = stream.next().await else { - vortex_bail!("flat layout received stream with more than a single chunk"); - }; - Ok(FlatLayout::new_with_metadata( - row_count, - stream.dtype().clone(), - segment_id, - ReadContext::new(ctx.array_ctx().to_ids()), - array_node, - ) - .into_layout()) + let segment_id = self.segment_sink.write(sequence_id, buffers).await?; + self.layout = Some( + FlatLayout::new_with_metadata( + row_count, + self.dtype.clone(), + segment_id, + ReadContext::new(self.ctx.array_ctx().to_ids()), + array_node, + ) + .into_layout(), + ); + Ok(()) + } + + async fn finish(&mut self, _sequence_id: SequenceId) -> VortexResult<()> { + Ok(()) + } + + async fn close(self: Box) -> VortexResult { + let Self { layout, dtype, .. } = *self; + Ok(layout.unwrap_or_else(|| { + ChunkedLayout::new(0, dtype, OwnedLayoutChildren::layout_children(vec![])).into_layout() + })) } } diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 4d8565fdd10..05106920039 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -4,8 +4,6 @@ use std::sync::Arc; use async_trait::async_trait; -use futures::StreamExt; -use futures::future::try_join; use futures::future::try_join_all; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -26,25 +24,18 @@ use vortex_array::scalar_fn::fns::operators::Operator; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_io::kanal_ext::KanalExt; use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriter; use crate::LayoutWriterContext; use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::layouts::list::ListLayout; use crate::segments::SegmentSinkRef; -use crate::sequence::SendableSequentialStream; use crate::sequence::SequenceId; -use crate::sequence::SequencePointer; -use crate::sequence::SequentialStream; -use crate::sequence::SequentialStreamAdapter; -use crate::sequence::SequentialStreamExt; - -/// Item carried on each child sub-stream: a sequenced, materialized chunk. -type ChildChunk = VortexResult<(SequenceId, ArrayRef)>; +use crate::strategy::LayoutWriterActor; /// Strategy for writing list-typed arrays, with a fallback for non-list dtypes. /// @@ -111,22 +102,16 @@ impl ListLayoutStrategy { } } -#[async_trait] impl LayoutStrategy for ListLayoutStrategy { - async fn write_stream( + fn new_writer( &self, ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, - stream: SendableSequentialStream, - mut eof: SequencePointer, + dtype: DType, session: &VortexSession, - ) -> VortexResult { - let dtype = stream.dtype().clone(); + ) -> VortexResult> { if !dtype.is_list() { - return self - .fallback - .write_stream(ctx, segment_sink, stream, eof, session) - .await; + return self.fallback.new_writer(ctx, segment_sink, dtype, session); } let is_nullable = dtype.is_nullable(); @@ -139,132 +124,133 @@ impl LayoutStrategy for ListLayoutStrategy { // so definsively widen. let offsets_dtype = DType::Primitive(PType::U64, Nullability::NonNullable); - // One bounded sub-stream per child: elements, offsets, and (when nullable) validity. - let (elements_tx, elements_rx) = kanal::bounded_async::(1); - let (offsets_tx, offsets_rx) = kanal::bounded_async::(1); - let (validity_tx, validity_rx) = if is_nullable { - let (tx, rx) = kanal::bounded_async::(1); - (Some(tx), Some(rx)) - } else { - (None, None) - }; - - // Transpose the list column into its child sub-streams and rebase offsets to global - // positions. Kept joined with the child writers below so producer errors surface rather - // than being hidden as an early channel close. - let fanout_fut = transpose_list_column( - stream, - session.clone(), - elements_tx, - offsets_tx, - validity_tx, - ); - - // Spawn a writer per child sub-stream, concurrently. - let handle = session.handle(); - let mut child_specs: Vec<( - DType, - Arc, - kanal::AsyncReceiver, - )> = vec![ - (element_dtype, Arc::clone(&self.elements), elements_rx), - (offsets_dtype, Arc::clone(&self.offsets), offsets_rx), + let mut child_specs = vec![ + (element_dtype, Arc::clone(&self.elements)), + (offsets_dtype, Arc::clone(&self.offsets)), ]; - if let Some(validity_rx) = validity_rx { + if is_nullable { child_specs.push(( DType::Bool(Nullability::NonNullable), Arc::clone(&self.validity), - validity_rx, )); } - let layout_futures: Vec<_> = child_specs + let buffered_bytes = ctx.buffered_bytes_tracker().clone(); + let handle = session.handle(); + let children = child_specs .into_iter() - .map(|(child_dtype, strategy, rx)| { - let child_stream = - SequentialStreamAdapter::new(child_dtype, rx.into_stream().boxed()).sendable(); - let child_eof = eof.split_off(); - let ctx = ctx.clone(); - let segment_sink = Arc::clone(&segment_sink); - let session = session.clone(); - handle.spawn_nested(move |h| async move { - let session = session.with_handle(h); - strategy - .write_stream(ctx, segment_sink, child_stream, child_eof, &session) - .await - }) + .map(|(child_dtype, strategy)| { + let writer = strategy.new_writer( + ctx.clone(), + Arc::clone(&segment_sink), + child_dtype, + session, + )?; + Ok(LayoutWriterActor::spawn( + writer, + buffered_bytes.clone(), + &handle, + )) }) - .collect(); - - let (_, layouts) = try_join(fanout_fut, try_join_all(layout_futures)).await?; - let mut layouts = layouts.into_iter(); - let elements_layout = layouts.next().vortex_expect("elements layout present"); - let offsets_layout = layouts.next().vortex_expect("offsets layout present"); - let validity_layout = - is_nullable.then(|| layouts.next().vortex_expect("validity layout present")); - - Ok(ListLayout::new(dtype, elements_layout, offsets_layout, validity_layout).into_layout()) + .collect::>>()?; + + Ok(Box::new(ListLayoutWriter { + dtype, + is_nullable, + children, + exec_ctx: session.create_execution_ctx(), + element_base: 0, + first: true, + saw_chunk: false, + })) } } -/// Transpose a list column into its `elements`, `offsets`, and (when present) `validity` child -/// sub-streams, rebasing each chunk's local `offsets` to global `u64` positions so the single -/// `offsets` child indexes into the concatenated `elements` child. -/// -/// `validity_tx` is `Some` exactly when the list is nullable. Errors surface to the caller, which -/// joins this against the child writers, rather than being hidden as an early channel close. -async fn transpose_list_column( - mut stream: SendableSequentialStream, - session: VortexSession, - elements_tx: kanal::AsyncSender, - offsets_tx: kanal::AsyncSender, - validity_tx: Option>, -) -> VortexResult<()> { - let mut exec_ctx = session.create_execution_ctx(); - let mut element_base: u64 = 0; - let mut first = true; - let mut saw_chunk = false; - while let Some(chunk) = stream.next().await { - let (sequence_id, array) = chunk?; - saw_chunk = true; - let mut sp = sequence_id.descend(); +struct ListLayoutWriter { + dtype: DType, + is_nullable: bool, + children: Vec, + exec_ctx: ExecutionCtx, + element_base: u64, + first: bool, + saw_chunk: bool, +} + +#[async_trait] +impl LayoutWriter for ListLayoutWriter { + async fn write(&mut self, sequence_id: SequenceId, array: ArrayRef) -> VortexResult<()> { + self.saw_chunk = true; let ListDataParts { elements, offsets, validity, .. - } = canonicalize_to_list_parts(array, &mut exec_ctx)?; + } = canonicalize_to_list_parts(array, &mut self.exec_ctx)?; let n_elements = elements.len() as u64; let row_count = offsets.len().saturating_sub(1); - let offsets = global_offsets(offsets, element_base, first, &mut exec_ctx)?; - element_base += n_elements; - first = false; - - if elements_tx - .send(Ok((sp.advance(), elements))) - .await - .is_err() - || offsets_tx.send(Ok((sp.advance(), offsets))).await.is_err() - { - vortex_bail!("list child writer finished before all chunks were sent"); + let offsets = global_offsets(offsets, self.element_base, self.first, &mut self.exec_ctx)?; + self.element_base += n_elements; + self.first = false; + + let mut columns = vec![elements, offsets]; + if self.is_nullable { + columns.push( + validity + .execute_mask(row_count, &mut self.exec_ctx)? + .into_array(), + ); } - if let Some(validity_tx) = &validity_tx { - let validity = validity - .execute_mask(row_count, &mut exec_ctx)? - .into_array(); - if validity_tx - .send(Ok((sp.advance(), validity))) - .await - .is_err() - { - vortex_bail!("list validity writer finished before all chunks were sent"); - } + + let mut sequence = sequence_id.descend(); + let child_sequences = (0..self.children.len()) + .map(|_| sequence.advance()) + .collect::>(); + try_join_all( + self.children + .iter_mut() + .zip(columns) + .zip(child_sequences) + .map(|((writer, column), sequence_id)| writer.write(sequence_id, column)), + ) + .await?; + Ok(()) + } + + async fn finish(&mut self, sequence_id: SequenceId) -> VortexResult<()> { + if !self.saw_chunk { + vortex_bail!("ListLayoutStrategy needs at least one chunk"); } + let mut sequence = sequence_id.descend(); + let child_sequences = (0..self.children.len()) + .map(|_| sequence.advance()) + .collect::>(); + try_join_all( + self.children + .iter_mut() + .zip(child_sequences) + .map(|(child, sequence_id)| child.finish(sequence_id)), + ) + .await?; + Ok(()) } - if !saw_chunk { - vortex_bail!("ListLayoutStrategy needs at least one chunk"); + + async fn close(self: Box) -> VortexResult { + let Self { + dtype, + is_nullable, + children, + .. + } = *self; + let mut child_layouts = Vec::with_capacity(children.len()); + for mut writer in children { + child_layouts.push(writer.take_layout()?); + } + let mut layouts = child_layouts.into_iter(); + let elements = layouts.next().vortex_expect("elements layout present"); + let offsets = layouts.next().vortex_expect("offsets layout present"); + let validity = is_nullable.then(|| layouts.next().vortex_expect("validity layout present")); + Ok(ListLayout::new(dtype, elements, offsets, validity).into_layout()) } - Ok(()) } /// Canonicalize a list-dtype array into [`ListDataParts`]. @@ -322,6 +308,7 @@ impl Matcher for AnyList { #[cfg(test)] mod tests { + use futures::StreamExt; use futures::stream; use vortex_array::ArrayContext; use vortex_array::arrays::BoolArray; @@ -333,13 +320,17 @@ mod tests { use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_io::session::RuntimeSession; + use vortex_io::session::RuntimeSessionExt; use super::*; use crate::layouts::chunked::writer::ChunkedLayoutStrategy; use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::layouts::table::TableStrategy; use crate::segments::TestSegments; + use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; + use crate::sequence::SequentialStreamAdapter; + use crate::sequence::SequentialStreamExt; use crate::session::LayoutSession; fn layout_test_session() -> VortexSession { @@ -478,9 +469,9 @@ mod tests { insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list({a=i32, b=i32}), children: 2 ├── elements: vortex.struct, dtype: {a=i32, b=i32}, children: 2 - │ ├── a: vortex.flat, dtype: i32, segment: 1 - │ └── b: vortex.flat, dtype: i32, segment: 2 - └── offsets: vortex.flat, dtype: u64, segment: 0 + │ ├── a: vortex.flat, dtype: i32, segment: 0 + │ └── b: vortex.flat, dtype: i32, segment: 1 + └── offsets: vortex.flat, dtype: u64, segment: 2 "); Ok(()) } @@ -506,9 +497,9 @@ mod tests { insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list(list(i32)), children: 2 ├── elements: vortex.list, dtype: list(i32), children: 2 - │ ├── elements: vortex.flat, dtype: i32, segment: 1 - │ └── offsets: vortex.flat, dtype: u64, segment: 2 - └── offsets: vortex.flat, dtype: u64, segment: 0 + │ ├── elements: vortex.flat, dtype: i32, segment: 0 + │ └── offsets: vortex.flat, dtype: u64, segment: 1 + └── offsets: vortex.flat, dtype: u64, segment: 2 "); Ok(()) } @@ -539,10 +530,10 @@ mod tests { vortex.list, dtype: list(list(list(i32))), children: 2 ├── elements: vortex.list, dtype: list(list(i32)), children: 2 │ ├── elements: vortex.list, dtype: list(i32), children: 2 - │ │ ├── elements: vortex.flat, dtype: i32, segment: 2 - │ │ └── offsets: vortex.flat, dtype: u64, segment: 3 - │ └── offsets: vortex.flat, dtype: u64, segment: 1 - └── offsets: vortex.flat, dtype: u64, segment: 0 + │ │ ├── elements: vortex.flat, dtype: i32, segment: 0 + │ │ └── offsets: vortex.flat, dtype: u64, segment: 1 + │ └── offsets: vortex.flat, dtype: u64, segment: 2 + └── offsets: vortex.flat, dtype: u64, segment: 3 "); Ok(()) } diff --git a/vortex-layout/src/layouts/repartition.rs b/vortex-layout/src/layouts/repartition.rs index 7344ee7be51..2cd085fceaa 100644 --- a/vortex-layout/src/layouts/repartition.rs +++ b/vortex-layout/src/layouts/repartition.rs @@ -4,10 +4,7 @@ use std::collections::VecDeque; use std::sync::Arc; -use async_stream::try_stream; use async_trait::async_trait; -use futures::StreamExt as _; -use futures::pin_mut; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; @@ -18,14 +15,14 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; +use crate::BufferedBytesReservation; +use crate::BufferedBytesTracker; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriter; use crate::LayoutWriterContext; use crate::segments::SegmentSinkRef; -use crate::sequence::SendableSequentialStream; -use crate::sequence::SequencePointer; -use crate::sequence::SequentialStreamAdapter; -use crate::sequence::SequentialStreamExt; +use crate::sequence::SequenceId; #[derive(Clone)] pub struct RepartitionWriterOptions { @@ -62,7 +59,7 @@ impl RepartitionWriterOptions { match dtype.element_size() { Some(elem_size) if elem_size > 0 => { // `div_ceil` ensures we overshoot the block_size_target; therefore preventing - // `write_stream` from combining adjacent 0.9 MiB chunks into one 1.8 MiB chunk. + // Prevent adjacent 0.9 MiB inputs from being combined into one 1.8 MiB block. let max_rows = usize::try_from(block_size_target.div_ceil(elem_size as u64)) .unwrap_or(usize::MAX); self.block_len_multiple.min(max_rows).max(1) @@ -91,100 +88,111 @@ impl RepartitionStrategy { } } -#[async_trait] impl LayoutStrategy for RepartitionStrategy { - async fn write_stream( + fn new_writer( &self, ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, - stream: SendableSequentialStream, - eof: SequencePointer, + dtype: DType, session: &VortexSession, - ) -> VortexResult { - // TODO(os): spawn stream below like: - // canon_stream = stream.map(async {to_canonical}).map(spawn).buffered(parallelism) - let dtype = stream.dtype().clone(); - let stream = if self.options.canonicalize { - let canonicalize_session = session.clone(); - SequentialStreamAdapter::new( - dtype.clone(), - stream.map(move |chunk| { - let (sequence_id, chunk) = chunk?; - let mut ctx = canonicalize_session.create_execution_ctx(); - let canonical = chunk.execute::(&mut ctx)?.into_array(); - VortexResult::Ok((sequence_id, canonical)) - }), - ) - .sendable() + ) -> VortexResult> { + let block_len = self.options.effective_block_len(&dtype); + let buffered_bytes = ctx.buffered_bytes_tracker().clone(); + Ok(Box::new(RepartitionLayoutWriter { + child: self + .child + .new_writer(ctx, segment_sink, dtype.clone(), session)?, + chunks: ChunksBuffer::new(self.options.block_size_minimum, block_len, buffered_bytes), + dtype, + block_len, + canonicalize: self.options.canonicalize, + exec_ctx: session.create_execution_ctx(), + })) + } +} + +struct RepartitionLayoutWriter { + child: Box, + chunks: ChunksBuffer, + dtype: DType, + block_len: usize, + canonicalize: bool, + exec_ctx: vortex_array::ExecutionCtx, +} + +impl RepartitionLayoutWriter { + fn canonicalize(&mut self, chunk: ArrayRef) -> VortexResult { + if self.canonicalize { + Ok(chunk.execute::(&mut self.exec_ctx)?.into_array()) } else { - stream - }; + Ok(chunk) + } + } - let dtype_clone = dtype.clone(); - let options = self.options.clone(); - - // For fixed-width types with large per-element sizes, reduce the block_len_multiple - // so that each block targets block_size_target bytes rather than producing oversized - // segments. - let block_len = options.effective_block_len(&dtype); - let block_size_minimum = options.block_size_minimum; - let repartition_session = session.clone(); - - let repartitioned_stream = try_stream! { - let canonical_stream = stream.peekable(); - pin_mut!(canonical_stream); - - let mut ctx = repartition_session.create_execution_ctx(); - let mut chunks = ChunksBuffer::new(block_size_minimum, block_len); - while let Some(chunk) = canonical_stream.as_mut().next().await { - let (sequence_id, chunk) = chunk?; - let mut sequence_pointer = sequence_id.descend(); - let mut offset = 0; - while offset < chunk.len() { - let end = (offset + block_len).min(chunk.len()); - let sliced = chunk.slice(offset..end)?; - chunks.push_back(sliced); - offset = end; - - if chunks.have_enough() { - let output_chunks = chunks.collect_exact_blocks()?; - assert!(!output_chunks.is_empty()); - let chunked = - ChunkedArray::try_new(output_chunks, dtype_clone.clone())?; - if !chunked.is_empty() { - let canonical = chunked.into_array().execute::(&mut ctx)?.into_array(); - yield ( - sequence_pointer.advance(), - canonical, - ) - } - } - } - if canonical_stream.as_mut().peek().await.is_none() { - let to_flush = ChunkedArray::try_new( - chunks.data.drain(..).map(|(arr, _)| arr), - dtype_clone.clone(), - )?; - if !to_flush.is_empty() { - let canonical = to_flush.into_array().execute::(&mut ctx)?.into_array(); - yield ( - sequence_pointer.advance(), - canonical, - ) - } + fn build_output( + &mut self, + mut chunks: Vec<(SequenceId, ArrayRef)>, + ) -> VortexResult<(SequenceId, ArrayRef)> { + let (sequence_id, last) = chunks.pop().vortex_expect("output chunks are non-empty"); + let chunked = ChunkedArray::try_new( + chunks + .into_iter() + .map(|(_, chunk)| chunk) + .chain(std::iter::once(last)), + self.dtype.clone(), + )?; + Ok(( + sequence_id, + chunked + .into_array() + .execute::(&mut self.exec_ctx)? + .into_array(), + )) + } +} + +#[async_trait] +impl LayoutWriter for RepartitionLayoutWriter { + async fn write(&mut self, sequence_id: SequenceId, chunk: ArrayRef) -> VortexResult<()> { + let chunk = self.canonicalize(chunk)?; + let mut sequence = sequence_id.descend(); + let mut offset = 0; + while offset < chunk.len() { + let end = (offset + self.block_len).min(chunk.len()); + self.chunks + .push_back(sequence.advance(), chunk.slice(offset..end)?); + offset = end; + + if self.chunks.have_enough() { + let chunks = self.chunks.collect_exact_blocks()?; + assert!(!chunks.is_empty()); + let (sequence_id, output) = self.build_output(chunks)?; + if !output.is_empty() { + self.child.write(sequence_id, output).await?; } } - }; + } + Ok(()) + } - self.child - .write_stream( - ctx, - segment_sink, - SequentialStreamAdapter::new(dtype, repartitioned_stream).sendable(), - eof, - session, - ) - .await + async fn finish(&mut self, sequence_id: SequenceId) -> VortexResult<()> { + if !self.chunks.data.is_empty() { + let chunks = self + .chunks + .data + .drain(..) + .map(|(sequence_id, array, _, _reservation)| (sequence_id, array)) + .collect(); + let (sequence_id, output) = self.build_output(chunks)?; + if !output.is_empty() { + self.child.write(sequence_id, output).await?; + } + } + self.child.finish(sequence_id).await + } + + async fn close(self: Box) -> VortexResult { + self.child.close().await } } @@ -192,21 +200,27 @@ struct ChunksBuffer { /// Each entry stores the chunk and the `nbytes()` snapshot taken at push time. /// This avoids accounting mismatches when interior-mutable arrays (e.g. `SharedArray`) /// change their reported size after being pushed. - data: VecDeque<(ArrayRef, u64)>, + data: VecDeque<(SequenceId, ArrayRef, u64, BufferedBytesReservation)>, row_count: usize, nbytes: u64, block_size_minimum: u64, block_len_multiple: usize, + buffered_bytes: BufferedBytesTracker, } impl ChunksBuffer { - fn new(block_size_minimum: u64, block_len_multiple: usize) -> Self { + fn new( + block_size_minimum: u64, + block_len_multiple: usize, + buffered_bytes: BufferedBytesTracker, + ) -> Self { Self { data: Default::default(), row_count: 0, nbytes: 0, block_size_minimum, block_len_multiple, + buffered_bytes, } } @@ -214,12 +228,12 @@ impl ChunksBuffer { self.nbytes >= self.block_size_minimum && self.row_count >= self.block_len_multiple } - fn collect_exact_blocks(&mut self) -> VortexResult> { + fn collect_exact_blocks(&mut self) -> VortexResult> { let nblocks = self.row_count / self.block_len_multiple; let mut res = Vec::with_capacity(self.data.len()); let mut remaining = nblocks * self.block_len_multiple; while remaining > 0 { - let (chunk, _) = self + let (sequence_id, chunk, _, _reservation) = self .pop_front() .vortex_expect("must have at least one chunk"); let len = chunk.len(); @@ -227,34 +241,37 @@ impl ChunksBuffer { if len > remaining { let left = chunk.slice(0..remaining)?; let right = chunk.slice(remaining..len)?; - self.push_front(right); - res.push(left); + let mut sequence = sequence_id.descend(); + res.push((sequence.advance(), left)); + self.push_front(sequence.advance(), right); remaining = 0; } else { - res.push(chunk); + res.push((sequence_id, chunk)); remaining -= len; } } Ok(res) } - fn push_back(&mut self, chunk: ArrayRef) { + fn push_back(&mut self, sequence_id: SequenceId, chunk: ArrayRef) { let nb = chunk.nbytes(); self.row_count += chunk.len(); self.nbytes += nb; - self.data.push_back((chunk, nb)); + self.data + .push_back((sequence_id, chunk, nb, self.buffered_bytes.reserve(nb))); } - fn push_front(&mut self, chunk: ArrayRef) { + fn push_front(&mut self, sequence_id: SequenceId, chunk: ArrayRef) { let nb = chunk.nbytes(); self.row_count += chunk.len(); self.nbytes += nb; - self.data.push_front((chunk, nb)); + self.data + .push_front((sequence_id, chunk, nb, self.buffered_bytes.reserve(nb))); } - fn pop_front(&mut self) -> Option<(ArrayRef, u64)> { + fn pop_front(&mut self) -> Option<(SequenceId, ArrayRef, u64, BufferedBytesReservation)> { let res = self.data.pop_front(); - if let Some((chunk, nb)) = res.as_ref() { + if let Some((_, chunk, nb, _)) = res.as_ref() { self.row_count -= chunk.len(); self.nbytes -= nb; } @@ -508,11 +525,16 @@ mod tests { let s1 = arr.slice(0..block_len)?; let s2 = arr.slice(block_len..n)?; - let mut buf = ChunksBuffer::new(0, block_len); - buf.push_back(s1); - buf.push_back(s2); + let tracker = BufferedBytesTracker::new(); + let mut buf = ChunksBuffer::new(0, block_len, tracker.clone()); + let mut sequence = SequenceId::root(); + buf.push_back(sequence.advance(), s1); + buf.push_back(sequence.advance(), s2); + + assert!(tracker.buffered_bytes() > 0); - let _output = buf.pop_front().unwrap(); + let output = buf.pop_front().unwrap(); + drop(output); // Transition SharedState from Source to Cached for ALL slices sharing this Arc. use vortex_array::arrays::shared::SharedArrayExt; @@ -520,9 +542,11 @@ mod tests { shared_handle.get_or_compute(|source| source.clone().execute::(&mut ctx))?; // Before the fix this panicked with "attempt to subtract with overflow". - let _s2 = buf.pop_front().unwrap(); + let s2 = buf.pop_front().unwrap(); + drop(s2); assert_eq!(buf.nbytes, 0); assert_eq!(buf.row_count, 0); + assert_eq!(tracker.buffered_bytes(), 0); Ok(()) } diff --git a/vortex-layout/src/layouts/struct_/writer.rs b/vortex-layout/src/layouts/struct_/writer.rs index fee71596fee..c4f74d54ab7 100644 --- a/vortex-layout/src/layouts/struct_/writer.rs +++ b/vortex-layout/src/layouts/struct_/writer.rs @@ -14,11 +14,7 @@ use std::sync::Arc; use async_trait::async_trait; -use futures::StreamExt; -use futures::TryStreamExt; -use futures::future::try_join; use futures::future::try_join_all; -use futures::pin_mut; use itertools::Itertools; use vortex_array::ArrayRef; use vortex_array::IntoArray; @@ -28,10 +24,8 @@ use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::FieldName; use vortex_array::dtype::Nullability; -use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_io::kanal_ext::KanalExt; use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; use vortex_utils::aliases::DefaultHashBuilder; @@ -40,14 +34,12 @@ use vortex_utils::aliases::hash_set::HashSet; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriter; use crate::LayoutWriterContext; use crate::layouts::struct_::StructLayout; use crate::segments::SegmentSinkRef; -use crate::sequence::SendableSequentialStream; use crate::sequence::SequenceId; -use crate::sequence::SequencePointer; -use crate::sequence::SequentialStreamAdapter; -use crate::sequence::SequentialStreamExt; +use crate::strategy::LayoutWriterActor; /// Writes struct-typed arrays into a [`StructLayout`], one child layout per field. /// @@ -101,18 +93,14 @@ impl StructStrategy { } } -#[async_trait] impl LayoutStrategy for StructStrategy { - async fn write_stream( + fn new_writer( &self, ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, - stream: SendableSequentialStream, - mut eof: SequencePointer, + dtype: DType, session: &VortexSession, - ) -> VortexResult { - let dtype = stream.dtype().clone(); - + ) -> VortexResult> { let Some(struct_dtype) = dtype.as_struct_fields_opt() else { vortex_bail!("StructStrategy can only write struct-typed streams, got {dtype}"); }; @@ -125,81 +113,6 @@ impl LayoutStrategy for StructStrategy { } let is_nullable = dtype.is_nullable(); - // Optimization: when there are no fields, don't spawn any work and just write a trivial - // StructLayout. - if struct_dtype.nfields() == 0 && !is_nullable { - let row_count = stream - .try_fold( - 0u64, - |acc, (_, arr)| async move { Ok(acc + arr.len() as u64) }, - ) - .await?; - return Ok(StructLayout::new(row_count, dtype, vec![]).into_layout()); - } - - // stream -> stream> - let columns_session = session.clone(); - let columns_vec_stream = stream.map(move |chunk| { - let (sequence_id, chunk) = chunk?; - let mut sequence_pointer = sequence_id.descend(); - let mut ctx = columns_session.create_execution_ctx(); - let struct_chunk = chunk.clone().execute::(&mut ctx)?; - let mut columns: Vec<(SequenceId, ArrayRef)> = Vec::new(); - if is_nullable { - columns.push(( - sequence_pointer.advance(), - chunk - .validity()? - .execute_mask(chunk.len(), &mut ctx)? - .into_array(), - )); - } - - columns.extend( - struct_chunk - .iter_unmasked_fields() - .map(|field| (sequence_pointer.advance(), field.clone())), - ); - - Ok(columns) - }); - - let mut stream_count = struct_dtype.nfields(); - if is_nullable { - stream_count += 1; - } - - let (column_streams_tx, column_streams_rx): (Vec<_>, Vec<_>) = - (0..stream_count).map(|_| kanal::bounded_async(1)).unzip(); - - // Fan out column chunks to their respective transposed streams. Keep this future joined - // with the column writers so producer panics/errors cannot be hidden as channel EOF. - let handle = session.handle(); - let fanout_fut = async move { - pin_mut!(columns_vec_stream); - while let Some(result) = columns_vec_stream.next().await { - match result { - Ok(columns) => { - for (tx, column) in column_streams_tx.iter().zip_eq(columns) { - if tx.send(Ok(column)).await.is_err() { - vortex_bail!( - "struct column writer finished before all chunks were sent" - ); - } - } - } - Err(e) => { - let e: Arc = Arc::new(e); - for tx in column_streams_tx.iter() { - let _ = tx.send(Err(VortexError::from(Arc::clone(&e)))).await; - } - return Err(VortexError::from(e)); - } - } - } - Ok(()) - }; - // First child column is the validity, subsequent children are the individual struct fields let column_dtypes: Vec = if is_nullable { std::iter::once(DType::Bool(Nullability::NonNullable)) @@ -217,44 +130,110 @@ impl LayoutStrategy for StructStrategy { struct_dtype.names().iter().cloned().collect() }; - let layout_futures: Vec<_> = column_dtypes + let buffered_bytes = ctx.buffered_bytes_tracker().clone(); + let handle = session.handle(); + let children = column_dtypes .into_iter() - .zip_eq(column_streams_rx) .zip_eq(column_names) .enumerate() - .map(move |(index, ((dtype, recv), name))| { - let column_stream = - SequentialStreamAdapter::new(dtype, recv.into_stream().boxed()).sendable(); - let child_eof = eof.split_off(); - let session = session.clone(); - let ctx = ctx.clone(); - let segment_sink = Arc::clone(&segment_sink); - handle.spawn_nested(move |h| { - // Validity is written through the validity strategy; every other field - // resolves to its named override or the default strategy. - let writer = if index == 0 && is_nullable { - Arc::clone(&self.validity) - } else { - self.field_writers - .get(&name) - .cloned() - .unwrap_or_else(|| Arc::clone(&self.default)) - }; - let session = session.with_handle(h); - - async move { - writer - .write_stream(ctx, segment_sink, column_stream, child_eof, &session) - .await - } - }) + .map(|(index, (dtype, name))| { + let strategy = if index == 0 && is_nullable { + Arc::clone(&self.validity) + } else { + self.field_writers + .get(&name) + .cloned() + .unwrap_or_else(|| Arc::clone(&self.default)) + }; + let writer = + strategy.new_writer(ctx.clone(), Arc::clone(&segment_sink), dtype, session)?; + Ok(LayoutWriterActor::spawn( + writer, + buffered_bytes.clone(), + &handle, + )) }) - .collect(); + .collect::>>()?; + + Ok(Box::new(StructLayoutWriter { + dtype, + is_nullable, + children, + exec_ctx: session.create_execution_ctx(), + row_count: 0, + })) + } +} + +struct StructLayoutWriter { + dtype: DType, + is_nullable: bool, + children: Vec, + exec_ctx: vortex_array::ExecutionCtx, + row_count: u64, +} + +#[async_trait] +impl LayoutWriter for StructLayoutWriter { + async fn write(&mut self, sequence_id: SequenceId, chunk: ArrayRef) -> VortexResult<()> { + self.row_count += chunk.len() as u64; + if self.children.is_empty() { + return Ok(()); + } - let (_success, column_layouts) = try_join(fanout_fut, try_join_all(layout_futures)).await?; - // TODO(os): transposed stream could count row counts as well, - // This must hold though, all columns must have the same row count of the struct layout - let row_count = column_layouts.first().map(|l| l.row_count()).unwrap_or(0); - Ok(StructLayout::new(row_count, dtype, column_layouts).into_layout()) + let struct_chunk = chunk.clone().execute::(&mut self.exec_ctx)?; + let mut columns = Vec::with_capacity(self.children.len()); + if self.is_nullable { + columns.push( + chunk + .validity()? + .execute_mask(chunk.len(), &mut self.exec_ctx)? + .into_array(), + ); + } + columns.extend(struct_chunk.iter_unmasked_fields().cloned()); + + let mut sequence = sequence_id.descend(); + let child_sequences = (0..self.children.len()) + .map(|_| sequence.advance()) + .collect::>(); + try_join_all( + self.children + .iter_mut() + .zip_eq(columns) + .zip(child_sequences) + .map(|((writer, column), sequence_id)| writer.write(sequence_id, column)), + ) + .await?; + Ok(()) + } + + async fn finish(&mut self, sequence_id: SequenceId) -> VortexResult<()> { + let mut sequence = sequence_id.descend(); + let child_sequences = (0..self.children.len()) + .map(|_| sequence.advance()) + .collect::>(); + try_join_all( + self.children + .iter_mut() + .zip(child_sequences) + .map(|(child, sequence_id)| child.finish(sequence_id)), + ) + .await?; + Ok(()) + } + + async fn close(self: Box) -> VortexResult { + let Self { + dtype, + children, + row_count, + .. + } = *self; + let mut layouts = Vec::with_capacity(children.len()); + for mut writer in children { + layouts.push(writer.take_layout()?); + } + Ok(StructLayout::new(row_count, dtype, layouts).into_layout()) } } diff --git a/vortex-layout/src/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 1a3c1adc524..a43b30a5083 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -3,7 +3,7 @@ //! A configurable writer strategy for tabular data. //! -//! [`TableStrategy`] is a *dispatcher*: it inspects the dtype of the stream it is handed and +//! [`TableStrategy`] is a *dispatcher*: it inspects the dtype of the writer it constructs and //! routes struct columns to [`StructStrategy`], list columns to [`ListLayoutStrategy`], and //! everything else to the configured leaf strategy. Because it hands *itself* (suitably descended) //! to those structural writers as the strategy for their children, arbitrarily nested struct/list @@ -16,7 +16,7 @@ use std::env; use std::sync::Arc; use std::sync::LazyLock; -use async_trait::async_trait; +use vortex_array::dtype::DType; use vortex_array::dtype::Field; use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldPath; @@ -25,14 +25,12 @@ use vortex_session::VortexSession; use vortex_utils::aliases::hash_map::HashMap; use vortex_utils::aliases::hash_set::HashSet; -use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriter; use crate::LayoutWriterContext; use crate::layouts::list::writer::ListLayoutStrategy; use crate::layouts::struct_::StructStrategy; use crate::segments::SegmentSinkRef; -use crate::sequence::SendableSequentialStream; -use crate::sequence::SequencePointer; /// Whether [`TableStrategy`] writes list fields using a [`ListLayoutStrategy`] by /// default. Disabled unless the environment variable `VORTEX_EXPERIMENTAL_LIST_LAYOUT` @@ -47,10 +45,10 @@ pub fn use_experimental_list_layout() -> bool { type ListLayoutFactory = Arc Arc + Send + Sync>; -/// A configurable strategy for writing nested tabular data, dispatching each (sub)stream to the +/// A configurable strategy for writing nested tabular data, dispatching each writer node to the /// structural writer for its dtype. /// -/// Dispatch rules, applied to the dtype of the stream handed to [`write_stream`]: +/// Dispatch rules, applied to the dtype handed to [`LayoutStrategy::new_writer`]: /// - **struct** → [`StructStrategy`], with each field written by its override (if any) or by a /// descended copy of this dispatcher. /// - **list** → [`ListLayoutStrategy`], with `elements` written by a descended copy of this @@ -59,8 +57,6 @@ type ListLayoutFactory = Arc Arc VortexResult { - let dtype = stream.dtype().clone(); - + ) -> VortexResult> { if dtype.is_struct() { return self .struct_strategy() - .write_stream(ctx, segment_sink, stream, eof, session) - .await; + .new_writer(ctx, segment_sink, dtype, session); } if dtype.is_list() && let Some(list_strategy) = self.list_strategy() { - return list_strategy - .write_stream(ctx, segment_sink, stream, eof, session) - .await; + return list_strategy.new_writer(ctx, segment_sink, dtype, session); } // Leaf: hand off to the leaf strategy. - self.leaf - .write_stream(ctx, segment_sink, stream, eof, session) - .await + self.leaf.new_writer(ctx, segment_sink, dtype, session) } } @@ -432,12 +419,12 @@ mod tests { .into_array(); let layout = write(&flat_table().with_list_layout(), outer).await?; - insta::assert_snapshot!(layout.display_tree(), @r" + insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list(list(i32)), children: 2 ├── elements: vortex.list, dtype: list(i32), children: 2 - │ ├── elements: vortex.flat, dtype: i32, segment: 1 - │ └── offsets: vortex.flat, dtype: u64, segment: 2 - └── offsets: vortex.flat, dtype: u64, segment: 0 + │ ├── elements: vortex.flat, dtype: i32, segment: 0 + │ └── offsets: vortex.flat, dtype: u64, segment: 1 + └── offsets: vortex.flat, dtype: u64, segment: 2 "); Ok(()) } @@ -463,14 +450,14 @@ mod tests { let st = StructArray::from_fields([("items", items)].as_slice())?.into_array(); let layout = write(&flat_table().with_list_layout(), st).await?; - insta::assert_snapshot!(layout.display_tree(), @r" + insta::assert_snapshot!(layout.display_tree(), @" vortex.struct, dtype: {items=list({a=i32, b=i32})?}, children: 1 └── items: vortex.list, dtype: list({a=i32, b=i32})?, children: 3 ├── elements: vortex.struct, dtype: {a=i32, b=i32}, children: 2 - │ ├── a: vortex.flat, dtype: i32, segment: 2 - │ └── b: vortex.flat, dtype: i32, segment: 3 - ├── offsets: vortex.flat, dtype: u64, segment: 0 - └── validity: vortex.flat, dtype: bool, segment: 1 + │ ├── a: vortex.flat, dtype: i32, segment: 0 + │ └── b: vortex.flat, dtype: i32, segment: 1 + ├── offsets: vortex.flat, dtype: u64, segment: 2 + └── validity: vortex.flat, dtype: bool, segment: 3 "); Ok(()) } @@ -502,13 +489,13 @@ mod tests { ) .with_list_layout(); let layout = write(&dispatcher, chunked).await?; - insta::assert_snapshot!(layout.display_tree(), @r" + insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list(i32), children: 2 ├── elements: vortex.chunked, dtype: i32, children: 2 │ ├── [0]: vortex.flat, dtype: i32, segment: 0 - │ └── [1]: vortex.flat, dtype: i32, segment: 1 + │ └── [1]: vortex.flat, dtype: i32, segment: 2 └── offsets: vortex.chunked, dtype: u64, children: 2 - ├── [0]: vortex.flat, dtype: u64, segment: 2 + ├── [0]: vortex.flat, dtype: u64, segment: 1 └── [1]: vortex.flat, dtype: u64, segment: 3 "); Ok(()) @@ -601,13 +588,13 @@ mod tests { let chunked = ChunkedArray::try_new(vec![c0, c1], dtype)?.into_array(); let layout = write(&dispatcher, chunked).await?; - insta::assert_snapshot!(layout.display_tree(), @r" + insta::assert_snapshot!(layout.display_tree(), @" vortex.struct, dtype: {a=i32, b=i32}, children: 2 ├── a: vortex.chunked, dtype: i32, children: 2 │ ├── [0]: vortex.flat, dtype: i32, segment: 0 - │ └── [1]: vortex.flat, dtype: i32, segment: 1 + │ └── [1]: vortex.flat, dtype: i32, segment: 2 └── b: vortex.chunked, dtype: i32, children: 2 - ├── [0]: vortex.flat, dtype: i32, segment: 2 + ├── [0]: vortex.flat, dtype: i32, segment: 1 └── [1]: vortex.flat, dtype: i32, segment: 3 "); Ok(()) diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index 4151679e6c3..d0230c51ae5 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -3,12 +3,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::collections::VecDeque; use std::num::NonZeroUsize; use std::sync::Arc; use async_trait::async_trait; -use futures::StreamExt as _; -use parking_lot::Mutex; +use futures::FutureExt; +use futures::future::BoxFuture; +use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; @@ -25,7 +27,8 @@ use vortex_array::aggregate_fn::fns::nan_count::NanCount; use vortex_array::aggregate_fn::fns::null_count::NullCount; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::DType; -use vortex_error::VortexError; +use vortex_array::scalar::Scalar; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_io::session::RuntimeSessionExt; @@ -34,17 +37,15 @@ use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriter; use crate::LayoutWriterContext; use crate::layouts::zoned::AggregateStatsAccumulator; use crate::layouts::zoned::ZonedLayout; use crate::layouts::zoned::aggregate_partials; use crate::layouts::zoned::schema::default_bounded_stat_max_bytes; use crate::segments::SegmentSinkRef; -use crate::sequence::SendableSequentialStream; +use crate::sequence::SequenceId; use crate::sequence::SequencePointer; -use crate::sequence::SequentialArrayStreamExt; -use crate::sequence::SequentialStreamAdapter; -use crate::sequence::SequentialStreamExt; /// Configuration for building zoned layouts. /// @@ -94,106 +95,149 @@ impl ZonedStrategy { } } -#[async_trait] impl LayoutStrategy for ZonedStrategy { - async fn write_stream( + fn new_writer( &self, ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, - stream: SendableSequentialStream, - mut eof: SequencePointer, + dtype: DType, session: &VortexSession, - ) -> VortexResult { + ) -> VortexResult> { let aggregate_fns = self .options .aggregate_fns .clone() - .unwrap_or_else(|| default_zoned_aggregate_fns(stream.dtype(), session)); - let compute_session = session.clone(); - - let stats_accumulator = Arc::new(Mutex::new(AggregateStatsAccumulator::new( - stream.dtype(), - &aggregate_fns, - ))); + .unwrap_or_else(|| default_zoned_aggregate_fns(&dtype, session)); + let stats_accumulator = AggregateStatsAccumulator::new(&dtype, &aggregate_fns); + let aggregate_fns = stats_accumulator.aggregate_fns(); // The accumulator has dropped the aggregates this dtype cannot hold, leaving the ones // this write would record. An aggregate the context forbids fails the write, like a // forbidden array or layout: dropping it silently would leave a file that prunes worse // than the caller asked for, with nothing in the output saying so. - let aggregate_fns = stats_accumulator.lock().aggregate_fns(); for aggregate_fn in aggregate_fns.iter() { if !ctx.allows_aggregate(&aggregate_fn.id()) { vortex_bail!("Aggregate {} not permitted by ctx", aggregate_fn.id()); } } + let buffered_bytes = ctx.buffered_bytes_tracker().clone(); + let data = self + .child + .new_writer(ctx.clone(), Arc::clone(&segment_sink), dtype, session)?; + + Ok(Box::new(ZonedLayoutWriter { + data, + stats_strategy: Arc::clone(&self.stats), + ctx, + segment_sink, + session: session.clone(), + stats_accumulator, + aggregate_fns, + buffered_bytes, + concurrency: self.options.concurrency.get(), + block_size: self.options.block_size, + pending: VecDeque::new(), + stats_sequence: None, + })) + } +} + +type ZoneFuture = BoxFuture< + 'static, + VortexResult<( + SequenceId, + ArrayRef, + Vec, + crate::BufferedBytesReservation, + )>, +>; + +struct ZonedLayoutWriter { + data: Box, + stats_strategy: Arc, + ctx: LayoutWriterContext, + segment_sink: SegmentSinkRef, + session: VortexSession, + stats_accumulator: AggregateStatsAccumulator, + aggregate_fns: Arc<[AggregateFnRef]>, + buffered_bytes: crate::BufferedBytesTracker, + concurrency: usize, + block_size: NonZeroUsize, + pending: VecDeque, + stats_sequence: Option, +} + +impl ZonedLayoutWriter { + async fn drain_one(&mut self) -> VortexResult<()> { + let Some(future) = self.pending.pop_front() else { + return Ok(()); + }; + let (sequence_id, chunk, partials, reservation) = future.await?; + self.stats_accumulator.push_partials(partials)?; + drop(reservation); + self.data.write(sequence_id, chunk).await + } +} - let stream_dtype = stream.dtype().clone(); - let concurrency = self.options.concurrency.get(); - let stream = stream - .map(move |item| { - let aggregate_fns = Arc::clone(&aggregate_fns); - let session = compute_session.clone(); - session.handle().spawn_cpu(move || { - let (sequence_id, chunk) = item?; +#[async_trait] +impl LayoutWriter for ZonedLayoutWriter { + async fn write(&mut self, sequence_id: SequenceId, chunk: ArrayRef) -> VortexResult<()> { + let aggregate_fns = Arc::clone(&self.aggregate_fns); + let session = self.session.clone(); + let reservation = self.buffered_bytes.reserve(chunk.nbytes()); + self.pending.push_back( + self.session + .handle() + .spawn_cpu(move || { let partials = aggregate_partials( &chunk, &aggregate_fns, &mut session.create_execution_ctx(), )?; - Ok::<_, VortexError>((sequence_id, chunk, partials)) + Ok((sequence_id, chunk, partials, reservation)) }) - }) - .buffered(concurrency); - - // Accumulate zone stats in stream order so the auxiliary table stays aligned with the - // data child. - let stats_accumulator2 = Arc::clone(&stats_accumulator); - let stream = SequentialStreamAdapter::new( - stream_dtype, - stream.map(move |item| { - let (sequence_id, chunk, partials) = item?; - stats_accumulator2.lock().push_partials(partials)?; - Ok((sequence_id, chunk)) - }), - ) - .sendable(); + .boxed(), + ); + if self.pending.len() >= self.concurrency { + self.drain_one().await?; + } + Ok(()) + } - let block_size = self.options.block_size; + async fn finish(&mut self, sequence_id: SequenceId) -> VortexResult<()> { + while !self.pending.is_empty() { + self.drain_one().await?; + } + let mut sequence = sequence_id.descend(); + self.data.finish(sequence.advance()).await?; + self.stats_sequence = Some(sequence); + Ok(()) + } - // The eof used for the data child should appear _before_ our own stats tables. - let data_eof = eof.split_off(); - let data_layout = self - .child - .write_stream( - ctx.clone(), - Arc::clone(&segment_sink), - stream, - data_eof, - session, - ) - .await?; - - let mut exec_ctx = session.create_execution_ctx(); - let Some((stats_array, aggregate_fns)) = - stats_accumulator.lock().as_array(&mut exec_ctx)? + async fn close(mut self: Box) -> VortexResult { + let data_layout = self.data.close().await?; + let Some((stats_array, aggregate_fns)) = self + .stats_accumulator + .as_array(&mut self.session.create_execution_ctx())? else { - // If we have no stats (e.g. the DType doesn't support them), then we just return the - // child layout. return Ok(data_layout); }; - // We must defer creating the stats table LayoutWriter until now, because the DType of - // the table depends on which stats were successfully computed. - let stats_stream = stats_array - .into_array() - .to_array_stream() - .sequenced(eof.split_off()); - let zones_layout = self - .stats - .write_stream(ctx, Arc::clone(&segment_sink), stats_stream, eof, session) - .await?; - + let stats_array = stats_array.into_array(); + let mut stats = self.stats_strategy.new_writer( + self.ctx, + self.segment_sink, + stats_array.dtype().clone(), + &self.session, + )?; + let mut stats_sequence = self + .stats_sequence + .take() + .vortex_expect("zoned writer must be finished before close"); + stats.write(stats_sequence.advance(), stats_array).await?; + stats.finish(stats_sequence.advance()).await?; + let zones_layout = stats.close().await?; Ok( - ZonedLayout::try_new(data_layout, zones_layout, block_size, aggregate_fns)? + ZonedLayout::try_new(data_layout, zones_layout, self.block_size, aggregate_fns)? .into_layout(), ) } diff --git a/vortex-layout/src/segments/shared.rs b/vortex-layout/src/segments/shared.rs index c794daf608e..93683d2c0ab 100644 --- a/vortex-layout/src/segments/shared.rs +++ b/vortex-layout/src/segments/shared.rs @@ -95,10 +95,9 @@ mod tests { // Add a segment to the test source let data = ByteBuffer::from(vec![1, 2, 3, 4]); - let seq_id = SequenceId::root().downgrade(); source .segments - .write(seq_id, vec![data.clone()]) + .write(SequenceId::root().downgrade(), vec![data.clone()]) .await .unwrap(); @@ -124,10 +123,9 @@ mod tests { // Add a segment let data = ByteBuffer::from(vec![5, 6, 7, 8]); - let seq_id = SequenceId::root().downgrade(); source .segments - .write(seq_id, vec![data.clone()]) + .write(SequenceId::root().downgrade(), vec![data.clone()]) .await .unwrap(); diff --git a/vortex-layout/src/segments/test.rs b/vortex-layout/src/segments/test.rs index d880d15cc1a..b27495aa227 100644 --- a/vortex-layout/src/segments/test.rs +++ b/vortex-layout/src/segments/test.rs @@ -41,9 +41,11 @@ impl SegmentSource for TestSegments { impl SegmentSink for TestSegments { async fn write( &self, - _sequence_id: SequenceId, + mut sequence_id: SequenceId, buffers: Vec, ) -> VortexResult { + sequence_id.collapse().await; + // Combine all the buffers since we're only a test implementation let mut buffer = ByteBufferMut::empty(); for segment in buffers { diff --git a/vortex-layout/src/strategy.rs b/vortex-layout/src/strategy.rs index 5a0b1025e4a..de6306c1255 100644 --- a/vortex-layout/src/strategy.rs +++ b/vortex-layout/src/strategy.rs @@ -9,25 +9,30 @@ use async_trait::async_trait; use futures::StreamExt; use vortex_array::ArrayContext; use vortex_array::ArrayId; +use vortex_array::ArrayRef; use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::dtype::DType; use vortex_array::normalize::NormalizeOptions; use vortex_array::normalize::Operation; use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_io::runtime::Handle; +use vortex_io::runtime::Task; use vortex_session::VortexSession; use vortex_utils::aliases::hash_set::HashSet; use crate::LayoutRef; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; -use crate::sequence::SequencePointer; -use crate::sequence::SequentialStreamAdapter; -use crate::sequence::SequentialStreamExt; +use crate::sequence::SequenceId; -/// A shared counter of the bytes that layout strategies are holding but have not yet emitted. +/// A shared counter of the logical byte size of arrays retained by layout strategies. /// /// Clones share the same counter, so a tracker can be handed to a writer before the write begins -/// and polled while it runs. Strategies report their own retained bytes with -/// [`Self::reserve`], which releases the reservation on drop. +/// and polled while it runs. This includes arrays queued for asynchronous strategy work, but not +/// allocator overhead, statistics-builder state, or buffering performed by the output sink. +/// Strategies report their own retained bytes with [`Self::reserve`], which releases the +/// reservation on drop. #[derive(Clone, Debug, Default)] pub struct BufferedBytesTracker(Arc); @@ -147,56 +152,174 @@ impl From for LayoutWriterContext { } } -/// Writes an ordered array stream into a layout tree and segment sink. +/// Creates a stateful writer node in a layout writer tree. /// /// Layout strategies are writer-side extension points. Strategies may repartition, buffer, /// collect columns, compute statistics, compress arrays, or delegate to child strategies before -/// finally emitting segments. They must preserve the logical row order represented by the -/// [`SequencePointer`]s in the input stream. +/// finally emitting segments. Each node receives arrays in logical row order. #[async_trait] pub trait LayoutStrategy: 'static + Send + Sync { - /// Asynchronously process an ordered stream of array chunks, emitting them into a sink and - /// returning the [`Layout`][crate::Layout] instance that can be parsed to retrieve the data - /// from rest. + /// Construct a writer for one dtype. /// - /// This trait uses the `#[async_trait]` attribute to denote that trait objects of this type - /// can be `Box`ed or `Arc`ed and shared around. Commonly, these strategies are composed to - /// form a operator of operations, each of which modifies the chunk stream in some way before - /// passing the data on to a downstream writer. - /// - /// # Sequencing and EOF - /// - /// The `stream` parameter is a stream of ordered array chunks, each of which is associated - /// with a sequence pointer that indicates its position in the overall array. By passing - /// around these pointers (essentially vector clocks), the writer can support concurrent - /// and parallel processing while maintaining a deterministic order of data in the file. /// The `ctx` parameter carries both array serialization state and writer-scoped accounting - /// through every child strategy. - /// - /// The `eof` parameter is a guaranteed to be greater than all sequence pointers in the stream. - /// - /// Because child strategies can write to the end-of-file pointer, it is very important that - /// **all strategies must await all children concurrently**. Otherwise it is possible to - /// deadlock if one child is waiting to write to EOF while your strategy is preventing the - /// stream from progressing to completion. - /// - /// # Blocking operations - /// - /// This is an async trait method, which will return a `BoxFuture` that you can await from - /// any runtime. Implementations should avoid directly performing blocking work within the - /// `write_stream`, and should instead spawn it onto an appropriate runtime or threadpool - /// dedicated to such work. + /// through every child strategy. Expensive work and independent children may run concurrently; + /// [`SequenceId`] preserves the logical segment order for sinks that require it. + fn new_writer( + &self, + ctx: LayoutWriterContext, + segment_sink: SegmentSinkRef, + dtype: DType, + session: &VortexSession, + ) -> VortexResult>; + + /// Drive this strategy from an existing stream. /// - /// Such operations are common, and include things like compression and parsing large blobs - /// of data, or serializing very large messages to flatbuffers. + /// This compatibility adapter is useful at stream-owning API boundaries and in tests. Layout + /// strategies compose by constructing child writers and calling [`LayoutWriter::write`] + /// directly. async fn write_stream( &self, ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, - stream: SendableSequentialStream, - eof: SequencePointer, + mut stream: SendableSequentialStream, + eof: crate::sequence::SequencePointer, session: &VortexSession, - ) -> VortexResult; + ) -> VortexResult { + let mut writer = self.new_writer(ctx, segment_sink, stream.dtype().clone(), session)?; + while let Some(chunk) = stream.next().await { + let (sequence_id, chunk) = chunk?; + writer.write(sequence_id, chunk).await?; + } + drop(stream); + writer.finish(eof.downgrade()).await?; + writer.close().await + } +} + +/// A stateful node in a push-based layout writer tree. +#[async_trait] +pub trait LayoutWriter: Send { + /// Push one ordered array into this node. + async fn write(&mut self, sequence_id: SequenceId, chunk: ArrayRef) -> VortexResult<()>; + + /// Establish the terminal input barrier and drain retained arrays and asynchronous work. + /// + /// Composite nodes call this on every child before closing any child. This lets strategies + /// such as zoned writers commit all primary data across sibling columns before emitting + /// derived metadata. This is called exactly once, and no arrays may be written afterward. + async fn finish(&mut self, sequence_id: SequenceId) -> VortexResult<()>; + + /// Consume this already-finished node and return the completed layout. + async fn close(self: Box) -> VortexResult; +} + +enum ActorMessage { + Write(SequenceId, ArrayRef, BufferedBytesReservation), + Finish(SequenceId), +} + +const CHILD_WRITER_QUEUE_CAPACITY: usize = 1; + +/// Drives one independent child writer on the runtime. Its bounded mailbox lets backpressure +/// from the segment sink propagate up to the public writer while retaining enough slack for +/// sibling writers to make progress independently. Queued arrays are accounted by +/// [`BufferedBytesTracker`]. +pub struct LayoutWriterActor { + sender: Option>, + task: Option>>, + layout: Option, + buffered_bytes: BufferedBytesTracker, +} + +impl LayoutWriterActor { + /// Spawn a writer with a bounded input mailbox on `handle`. + pub fn spawn( + mut writer: Box, + buffered_bytes: BufferedBytesTracker, + handle: &Handle, + ) -> Self { + let (sender, receiver) = kanal::bounded_async::(CHILD_WRITER_QUEUE_CAPACITY); + let task = handle.spawn(async move { + loop { + match receiver.recv().await { + Ok(ActorMessage::Write(sequence_id, chunk, reservation)) => { + writer.write(sequence_id, chunk).await?; + drop(reservation); + } + Ok(ActorMessage::Finish(sequence_id)) => { + writer.finish(sequence_id).await?; + return writer.close().await; + } + Err(_) => return Err(vortex_err!("layout child sender dropped before finish")), + } + } + }); + Self { + sender: Some(sender), + task: Some(task), + layout: None, + buffered_bytes, + } + } + + /// Push an array into the child, waiting when its mailbox is full. + pub async fn write(&mut self, sequence_id: SequenceId, chunk: ArrayRef) -> VortexResult<()> { + let reservation = self.buffered_bytes.reserve(chunk.nbytes()); + self.sender + .as_ref() + .ok_or_else(|| vortex_err!("layout child is already finished"))? + .send(ActorMessage::Write(sequence_id, chunk, reservation)) + .await + .map_err(|_| vortex_err!("layout child finished before all chunks were pushed")) + } + + /// Finish the child and wait for its layout to be produced. + pub async fn finish(&mut self, sequence_id: SequenceId) -> VortexResult<()> { + if self.layout.is_some() { + return Ok(()); + } + self.sender + .take() + .ok_or_else(|| vortex_err!("layout child sender is missing"))? + .send(ActorMessage::Finish(sequence_id)) + .await + .map_err(|_| vortex_err!("layout child finished before its terminal barrier"))?; + let task = self + .task + .take() + .ok_or_else(|| vortex_err!("layout child task is missing"))?; + self.layout = Some(task.await?); + Ok(()) + } + + /// Take the completed layout after [`Self::finish`]. + pub fn take_layout(&mut self) -> VortexResult { + self.layout + .take() + .ok_or_else(|| vortex_err!("layout child was not finished")) + } +} + +/// Drive a push-based layout writer from an existing sequential stream. +/// +/// This is an adapter for callers and tests that already own streams; strategies themselves are +/// composed exclusively through [`LayoutWriter::write`]. +pub async fn write_stream( + strategy: &dyn LayoutStrategy, + ctx: LayoutWriterContext, + segment_sink: SegmentSinkRef, + mut stream: SendableSequentialStream, + eof: crate::sequence::SequencePointer, + session: &VortexSession, +) -> VortexResult { + let mut writer = strategy.new_writer(ctx, segment_sink, stream.dtype().clone(), session)?; + while let Some(chunk) = stream.next().await { + let (sequence_id, chunk) = chunk?; + writer.write(sequence_id, chunk).await?; + } + drop(stream); + writer.finish(eof.downgrade()).await?; + writer.close().await } /// A layout strategy wrapper that rejects arrays containing encodings outside an allow-list. @@ -219,58 +342,111 @@ impl LayoutStrategyEncodingValidator { } } -#[async_trait] impl LayoutStrategy for LayoutStrategyEncodingValidator { - async fn write_stream( + fn new_writer( &self, ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, - stream: SendableSequentialStream, - eof: SequencePointer, + dtype: DType, session: &VortexSession, - ) -> VortexResult { - let dtype = stream.dtype().clone(); - let allowed_encodings = Arc::clone(&self.allowed_encodings); - let stream = stream.map(move |chunk| { - let (sequence_id, chunk) = chunk?; - let chunk = chunk.normalize(&mut NormalizeOptions { - allowed: &allowed_encodings, - operation: Operation::Error, - })?; - Ok((sequence_id, chunk)) - }); - - self.child - .write_stream( - ctx, - segment_sink, - SequentialStreamAdapter::new(dtype, stream).sendable(), - eof, - session, - ) - .await + ) -> VortexResult> { + Ok(Box::new(EncodingValidatorWriter { + child: self.child.new_writer(ctx, segment_sink, dtype, session)?, + allowed_encodings: Arc::clone(&self.allowed_encodings), + })) } } -#[async_trait] impl LayoutStrategy for Arc { - async fn write_stream( + fn new_writer( &self, ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, - stream: SendableSequentialStream, - eof: SequencePointer, + dtype: DType, session: &VortexSession, - ) -> VortexResult { - (**self) - .write_stream(ctx, segment_sink, stream, eof, session) - .await + ) -> VortexResult> { + (**self).new_writer(ctx, segment_sink, dtype, session) + } +} + +#[async_trait] +impl LayoutWriter for EncodingValidatorWriter { + async fn write(&mut self, sequence_id: SequenceId, chunk: ArrayRef) -> VortexResult<()> { + let chunk = chunk.normalize(&mut NormalizeOptions { + allowed: &self.allowed_encodings, + operation: Operation::Error, + })?; + self.child.write(sequence_id, chunk).await + } + + async fn finish(&mut self, sequence_id: SequenceId) -> VortexResult<()> { + self.child.finish(sequence_id).await + } + + async fn close(self: Box) -> VortexResult { + self.child.close().await } } +struct EncodingValidatorWriter { + child: Box, + allowed_encodings: Arc>, +} + #[cfg(test)] mod tests { + use std::sync::Arc; + + use async_trait::async_trait; + use futures::FutureExt; + use tokio::sync::Semaphore; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use vortex_io::session::RuntimeSessionExt; + + use crate::LayoutRef; + use crate::LayoutWriter; + use crate::children::OwnedLayoutChildren; + use crate::layouts::chunked::ChunkedLayout; + use crate::sequence::SequenceId; use crate::strategy::BufferedBytesTracker; + use crate::strategy::CHILD_WRITER_QUEUE_CAPACITY; + use crate::strategy::LayoutWriterActor; + use crate::test::new_session; + + struct BlockingWriter { + permits: Arc, + } + + #[async_trait] + impl LayoutWriter for BlockingWriter { + async fn write(&mut self, _sequence_id: SequenceId, _chunk: ArrayRef) -> VortexResult<()> { + self.permits + .acquire() + .await + .map_err(|_| vortex_err!("test semaphore closed"))? + .forget(); + Ok(()) + } + + async fn finish(&mut self, _sequence_id: SequenceId) -> VortexResult<()> { + Ok(()) + } + + async fn close(self: Box) -> VortexResult { + Ok(ChunkedLayout::new( + 0, + DType::Bool(Nullability::NonNullable), + OwnedLayoutChildren::layout_children(vec![]), + ) + .into_layout()) + } + } #[test] fn reservations_accumulate_and_release() { @@ -300,4 +476,43 @@ mod tests { drop(reservation); assert_eq!(observer.buffered_bytes(), 0); } + + #[tokio::test] + async fn child_writer_mailbox_applies_backpressure_when_full() -> VortexResult<()> { + let permits = Arc::new(Semaphore::new(0)); + let tracker = BufferedBytesTracker::new(); + let session = new_session().with_tokio(); + let mut actor = LayoutWriterActor::spawn( + Box::new(BlockingWriter { + permits: Arc::clone(&permits), + }), + tracker.clone(), + &session.handle(), + ); + let chunk = buffer![1u64].into_array(); + let chunk_bytes = chunk.nbytes(); + let mut sequence = SequenceId::root(); + + for _ in 0..=CHILD_WRITER_QUEUE_CAPACITY { + actor.write(sequence.advance(), chunk.clone()).await?; + } + assert_eq!( + tracker.buffered_bytes(), + (CHILD_WRITER_QUEUE_CAPACITY as u64 + 1) * chunk_bytes + ); + + assert!( + actor + .write(sequence.advance(), chunk) + .now_or_never() + .is_none(), + "a full mailbox must block the producer" + ); + + permits.add_permits(CHILD_WRITER_QUEUE_CAPACITY + 1); + actor.finish(sequence.advance()).await?; + actor.take_layout()?; + assert_eq!(tracker.buffered_bytes(), 0); + Ok(()) + } }