diff --git a/vortex-array/src/serde.rs b/vortex-array/src/serde.rs index 04291b9eef8..3e3c952e2a6 100644 --- a/vortex-array/src/serde.rs +++ b/vortex-array/src/serde.rs @@ -191,16 +191,12 @@ impl<'a> ArrayNodeFlatBuffer<'a> { &self, fbb: &mut FlatBufferBuilder<'fb>, ) -> VortexResult>> { - let encoding_idx = self - .ctx - .intern(&self.array.encoding_id()) - // TODO(ngates): write_flatbuffer should return a result if this can fail. - .ok_or_else(|| { - vortex_err!( - "Array encoding {} not permitted by ctx", - self.array.encoding_id() - ) - })?; + let encoding_idx = self.ctx.intern(&self.array.encoding_id()).ok_or_else(|| { + vortex_err!( + "Array encoding {} not permitted by ctx", + self.array.encoding_id() + ) + })?; let metadata_bytes = self.session.array_serialize(self.array)?.ok_or_else(|| { vortex_err!( diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index c97f7b737cb..16cef49cf0b 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -13,7 +13,6 @@ use async_trait::async_trait; use futures::FutureExt; use futures::StreamExt; use futures::future::BoxFuture; -use vortex::array::ArrayContext; use vortex::array::ArrayRef; use vortex::array::ArrayVTable; use vortex::array::MaskFuture; @@ -45,6 +44,7 @@ use vortex::layout::LayoutReader; use vortex::layout::LayoutReaderRef; use vortex::layout::LayoutRef; use vortex::layout::LayoutStrategy; +use vortex::layout::LayoutWriterContext; use vortex::layout::RowSplits; use vortex::layout::SplitRange; use vortex::layout::VTable; @@ -414,13 +414,12 @@ fn truncate_scalar_stat Option<(Scalar, bool)>>( impl LayoutStrategy for CudaFlatLayoutStrategy { async fn write_stream( &self, - ctx: ArrayContext, + ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, mut stream: SendableSequentialStream, _eof: SequencePointer, session: &VortexSession, ) -> VortexResult { - let ctx = ctx.clone(); let options = self.clone(); let Some(chunk) = stream.next().await else { vortex_bail!("CudaFlatLayoutStrategy needs a single chunk"); @@ -472,7 +471,7 @@ impl LayoutStrategy for CudaFlatLayoutStrategy { let host_buffers = extract_constant_buffers(&chunk); let buffers = chunk.serialize( - &ctx, + ctx.array_ctx(), session, &SerializeOptions { offset: 0, @@ -503,7 +502,7 @@ impl LayoutStrategy for CudaFlatLayoutStrategy { layout_children(Vec::new()), CudaFlatData { segment_id, - ctx: ReadContext::new(ctx.to_ids()), + ctx: ReadContext::new(ctx.array_ctx().to_ids()), array_tree, host_buffers: Arc::new(host_buffer_map), }, diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 69813ed8ce5..1bedc72b70a 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -52,6 +52,7 @@ use vortex_array::expr::select; use vortex_array::extension::datetime::TimeUnit; use vortex_array::extension::datetime::Timestamp; use vortex_array::extension::datetime::TimestampOptions; +use vortex_array::field_path; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::fns::pack::Pack; @@ -67,18 +68,25 @@ use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; +use vortex_edition::EditionSession; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_flatbuffers::footer as fb; use vortex_io::session::RuntimeSession; use vortex_layout::DynLayout; +use vortex_layout::LayoutStrategy; +use vortex_layout::layouts::buffered::BufferedStrategy; +use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy; use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; +use vortex_layout::layouts::struct_::StructStrategy; +use vortex_layout::layouts::table::TableStrategy; use vortex_layout::layouts::zoned::LegacyStats; use vortex_layout::layouts::zoned::Zoned; use vortex_layout::scan::scan_builder::ScanBuilder; use vortex_layout::scan::split_by::SplitBy; use vortex_layout::session::LayoutSession; use vortex_session::VortexSession; +use vortex_zigzag::ZigZag; use crate::MAX_POSTSCRIPT_SIZE; use crate::OpenOptionsSessionExt; @@ -1377,7 +1385,7 @@ async fn test_into_tokio_array_stream() -> VortexResult<()> { ]) .into_array(); - let st = StructArray::from_fields(&[("strings", strings), ("numbers", numbers)]).unwrap(); + let st = StructArray::from_fields(&[("strings", strings), ("numbers", numbers)])?; let mut buf = ByteBufferMut::empty(); SESSION .write_options() @@ -1542,6 +1550,143 @@ async fn test_writer_bytes_written() -> VortexResult<()> { Ok(()) } +#[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)] +#[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; + + let fields = [ + ("a", buffer![1u32, 2, 3, 4].into_array()), + ("b", buffer![5u32, 6, 7, 8].into_array()), + ]; + let array = StructArray::from_fields(&fields[..leaf_count])?.into_array(); + + let new_leaf = || -> Arc { + Arc::new(BufferedStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + BUFFER_SIZE, + )) + }; + let validity: Arc = Arc::new(FlatLayoutStrategy::default()); + let strategy: Arc = if use_table_strategy { + let mut table = TableStrategy::new(validity, new_leaf()); + if field_override { + table = table.with_field_writer(field_path!(b), new_leaf()); + } + Arc::new(table) + } else { + Arc::new(StructStrategy::new(validity, new_leaf())) + }; + + 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()); + + 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. + writer.push(array.clone()).await?; + writer.push(array.clone()).await?; + writer.push(array).await?; + + assert_eq!(writer.buffered_bytes(), expected_buffered_bytes); + + let summary = writer.finish().await?; + assert_eq!(summary.row_count(), 12); + assert_eq!(buffered_bytes.buffered_bytes(), 0); + + Ok(()) +} + +#[tokio::test] +async fn test_buffered_bytes_are_writer_scoped() -> VortexResult<()> { + const BUFFER_SIZE: u64 = 16; + + let array = + StructArray::from_fields(&[("a", buffer![1u32, 2, 3, 4].into_array())])?.into_array(); + let leaf = Arc::new(BufferedStrategy::new( + ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), + BUFFER_SIZE, + )); + let strategy: Arc = Arc::new(TableStrategy::new( + Arc::new(FlatLayoutStrategy::default()), + leaf, + )); + + let mut first_buf = ByteBufferMut::empty(); + let mut first = SESSION + .write_options() + .with_strategy(Arc::clone(&strategy)) + .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()); + + first.push(array.clone()).await?; + first.push(array.clone()).await?; + first.push(array.clone()).await?; + second.push(array.clone()).await?; + 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); + + first.finish().await?; + second.finish().await?; + + Ok(()) +} + +#[tokio::test] +async fn test_encoding_registered_after_write_options() -> VortexResult<()> { + // A session that does not know about ZigZag yet. + let session = array_session() + .with::() + .with::() + .with::(); + + // Configure the options before the encoding is registered; `write` is what snapshots the + // session's encodings, so registering in between must still be honoured. + let options = session + .write_options() + .with_strategy(Arc::new(FlatLayoutStrategy::default())); + vortex_zigzag::initialize(&session); + crate::enable_all_registered_array_encodings(&session); + + let array = ZigZag::try_new(buffer![1u32, 2, 3, 4].into_array())?.into_array(); + let dtype = array.dtype().clone(); + + let mut buf = ByteBufferMut::empty(); + options.write(&mut buf, array.to_array_stream()).await?; + + let chunks: Vec<_> = session + .open_options() + .open_buffer(buf)? + .scan()? + .into_array_stream()? + .try_collect() + .await?; + let read = ChunkedArray::try_new(chunks, dtype)?.into_array(); + let mut ctx = session.create_execution_ctx(); + assert_arrays_eq!(read, buffer![-1i32, 1, -2, 2].into_array(), &mut ctx); + + Ok(()) +} + #[tokio::test] async fn test_writer_empty_chunks() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); @@ -2029,7 +2174,7 @@ async fn timestamp_unit_mismatch_errors_with_constant_children() "Expected error from timestamp unit mismatch (ms vs s), but got {} results. \ This indicates the scanner silently applied the filter incorrectly when \ DateTimePartsArray children use ConstantArray encoding.", - results.unwrap().len() + results?.len() ); Ok(()) @@ -2223,7 +2368,7 @@ async fn test_segment_ordering_dict_codes_before_values() -> VortexResult<()> { let strings = VarBinArray::from(values).into_array(); let numbers = PrimitiveArray::from_iter(0..n as i32).into_array(); - let st = StructArray::from_fields(&[("strings", strings), ("numbers", numbers)]).unwrap(); + let st = StructArray::from_fields(&[("strings", strings), ("numbers", numbers)])?; let mut buf = ByteBufferMut::empty(); let summary = SESSION @@ -2345,8 +2490,7 @@ async fn test_segment_ordering_zonemaps_after_data() -> VortexResult<()> { ("strings", strings), ("numbers", numbers), ("floats", floats), - ]) - .unwrap(); + ])?; let mut buf = ByteBufferMut::empty(); let summary = SESSION diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs index fc14eaf57d7..141b6b8ca73 100644 --- a/vortex-file/src/v2/file_stats_reader.rs +++ b/vortex-file/src/v2/file_stats_reader.rs @@ -246,7 +246,7 @@ mod tests { ); let layout = strategy .write_stream( - ctx, + ctx.into(), Arc::::clone(&segments), struct_array.into_array().to_array_stream().sequenced(ptr), eof, @@ -285,7 +285,7 @@ mod tests { ); let layout = strategy .write_stream( - ctx, + ctx.into(), Arc::::clone(&segments), struct_array.into_array().to_array_stream().sequenced(ptr), eof, @@ -324,7 +324,7 @@ mod tests { ); let layout = strategy .write_stream( - ctx, + ctx.into(), Arc::::clone(&segments), struct_array.into_array().to_array_stream().sequenced(ptr), eof, @@ -373,7 +373,7 @@ mod tests { ); let layout = strategy .write_stream( - ctx, + ctx.into(), Arc::clone(&segments) as Arc, struct_array.into_array().to_array_stream().sequenced(ptr), eof, @@ -422,7 +422,7 @@ mod tests { ); let layout = strategy .write_stream( - ctx, + ctx.into(), Arc::clone(&segments) as Arc, struct_array.into_array().to_array_stream().sequenced(ptr), eof, diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 90889579bc6..162476b23eb 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -15,6 +15,7 @@ use futures::future::LocalBoxFuture; use futures::future::ready; use futures::pin_mut; use futures::select; +use itertools::Itertools; use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::dtype::DType; @@ -39,7 +40,9 @@ 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::LayoutStrategy; +use vortex_layout::LayoutWriterContext; use vortex_layout::layouts::file_stats::accumulate_stats; use vortex_layout::sequence::SequenceId; use vortex_layout::sequence::SequentialStreamAdapter; @@ -68,6 +71,7 @@ use crate::segments::writer::BufferedSegmentSink; pub struct VortexWriteOptions { session: VortexSession, strategy: Arc, + buffered_bytes: BufferedBytesTracker, exclude_dtype: bool, max_variable_length_statistics_size: usize, file_statistics: Vec, @@ -78,8 +82,7 @@ pub struct VortexWriteOptions { pub trait WriteOptionsSessionExt: SessionExt { /// Create [`VortexWriteOptions`] for writing to a Vortex file. fn write_options(&self) -> VortexWriteOptions { - let session = self.session(); - VortexWriteOptions::new(session) + VortexWriteOptions::new(self.session()) } } impl WriteOptionsSessionExt for S {} @@ -92,6 +95,7 @@ impl VortexWriteOptions { .build(); VortexWriteOptions { strategy, + buffered_bytes: BufferedBytesTracker::new(), session, exclude_dtype: false, file_statistics: PRUNING_STATS.to_vec(), @@ -110,6 +114,16 @@ impl VortexWriteOptions { self } + /// Returns the tracker accounting for bytes that layout strategies are holding but have not + /// yet emitted. + /// + /// The tracker is shared with the write these options start, so it can be captured before + /// calling [`Self::write`] and polled while the write runs. [`Writer::buffered_bytes`] exposes + /// the same counter for the push-based API. + pub fn buffered_bytes_tracker(&self) -> BufferedBytesTracker { + self.buffered_bytes.clone() + } + /// Exclude the DType from the Vortex file. You must provide the DType to the reader. // TODO(ngates): Should we store some sort of DType checksum to make sure the one passed at // read-time is sane? I guess most layouts will have some reasonable validation. @@ -180,6 +194,9 @@ impl VortexWriteOptions { /// /// Note that buffers are flushed as soon as they are available with no buffering, the caller /// is responsible for deciding how to configure buffering on the underlying `Write` sink. + /// + /// The set of encodings permitted in the file is snapshotted from the session's array registry + /// here, so encodings registered after this call are not written. pub async fn write( self, write: W, @@ -196,15 +213,10 @@ impl VortexWriteOptions { ) -> VortexResult { validate_metadata_segments(&self.metadata)?; - // NOTE(os): Setup an array context that already has all known encodings pre-populated. - // This is preferred for now over having an empty context here, because only the - // serialised array order is deterministic. The serialisation of arrays are done - // parallel and with an empty context they can register their encodings to the context - // in different order, changing the written bytes from run to run. - let enabled_encoding_ids = self.session.enabled_encoding_ids(); - let ctx = ArrayContext::new(enabled_encoding_ids.clone()) - // Only permit encodings in the session's enabled editions. - .with_allowed_ids(enabled_encoding_ids.into_iter().collect()); + // 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 ctx = LayoutWriterContext::new(new_array_context(&self.session)) + .with_buffered_bytes_tracker(self.buffered_bytes.clone()); let dtype = stream.dtype().clone(); let (mut ptr, eof) = SequenceId::root().split(); @@ -277,7 +289,7 @@ impl VortexWriteOptions { Arc::clone(&layout), segment_specs, statistics, - ReadContext::new(ctx.to_ids()), + ReadContext::new(ctx.array_ctx().to_ids()), ); // Emit the footer buffers and EOF. @@ -318,18 +330,30 @@ impl VortexWriteOptions { let write = CountingVortexWrite::new(write); let bytes_written = write.counter(); - let strategy = Arc::clone(&self.strategy); + let buffered_bytes = self.buffered_bytes.clone(); let future = self.write(write, arrays).boxed_local().fuse(); Writer { arrays: Some(arrays_send), future, bytes_written, - strategy, + buffered_bytes, } } } +fn new_array_context(session: &VortexSession) -> ArrayContext { + // NOTE(os): Setup an array context that already has all known encodings pre-populated. + // This is preferred for now over having an empty context here, because only the + // serialised array order is deterministic. The serialisation of arrays are done + // parallel and with an empty context they can register their encodings to the context + // in different order, changing the written bytes from run to run. + let enabled_encoding_ids = session.enabled_encoding_ids(); + ArrayContext::new(enabled_encoding_ids.iter().cloned().sorted().collect()) + // Only permit encodings known to the session. + .with_allowed_ids(enabled_encoding_ids.into_iter().collect()) +} + fn validate_metadata_segments(metadata: &HashMap) -> VortexResult<()> { if metadata.len() > MAX_METADATA_SEGMENTS { vortex_bail!( @@ -370,8 +394,8 @@ pub struct Writer<'w> { future: Fuse>>, // The bytes written so far. bytes_written: Arc, - // The layout strategy that is being used for the write. - strategy: Arc, + // The buffered bytes accounting shared with the layout strategies for this write. + buffered_bytes: BufferedBytesTracker, } impl Writer<'_> { @@ -450,7 +474,7 @@ impl Writer<'_> { /// Returns the number of bytes currently buffered by the layout writers. pub fn buffered_bytes(&self) -> u64 { - self.strategy.buffered_bytes() + self.buffered_bytes.buffered_bytes() } /// Finish writing the Vortex file, flushing any remaining buffers and returning the @@ -659,11 +683,11 @@ mod tests { #[case::oversized_multibyte_key( vec!["é".repeat(MAX_METADATA_KEY_BYTES / "é".len() + 1)], "keys must be at most" - )] + )] #[case::too_many_segments( (0..=MAX_METADATA_SEGMENTS).map(|idx| format!("key-{idx}")).collect(), "at most 16 metadata segments" - )] + )] fn validate_metadata_rejects(#[case] keys: Vec, #[case] expected: &str) { let Err(error) = write_options_with_keys(&keys).validate_metadata() else { panic!("metadata must be rejected for {keys:?}"); diff --git a/vortex-jni/src/writer.rs b/vortex-jni/src/writer.rs index 92ea22349cc..9ed175d7e60 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -59,6 +59,7 @@ use vortex::io::object_store::ObjectStoreWrite; use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::Task; use vortex::io::session::RuntimeSessionExt; +use vortex::layout::BufferedBytesTracker; use vortex::layout::LayoutStrategy; use vortex::session::VortexSession; use vortex::utils::aliases::hash_map::HashMap; @@ -143,7 +144,7 @@ pub struct NativeWriter { arrow_schema: SchemaRef, write_schema: DType, bytes_written: Arc, - strategy: Arc, + buffered_bytes: BufferedBytesTracker, sender: mpsc::Sender>, } @@ -153,7 +154,7 @@ impl NativeWriter { arrow_schema: SchemaRef, write_schema: DType, bytes_written: Arc, - strategy: Arc, + buffered_bytes: BufferedBytesTracker, handle: Task>, sender: mpsc::Sender>, ) -> Self { @@ -163,7 +164,7 @@ impl NativeWriter { arrow_schema, write_schema, bytes_written, - strategy, + buffered_bytes, sender, } } @@ -206,7 +207,7 @@ impl NativeWriter { } fn buffered_bytes(&self) -> u64 { - self.strategy.buffered_bytes() + self.buffered_bytes.buffered_bytes() } fn close(mut self) -> VortexResult { @@ -424,11 +425,12 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create( let strategy = write_strategy_for_schema(session, &write_schema); let write_options = session .write_options() - .with_strategy(Arc::clone(&strategy)) + .with_strategy(strategy) .with_metadata_segments(metadata); // The same check runs inside `write`, but only once the write task is under way, where // it would surface as an opaque send failure on the first batch. write_options.validate_metadata()?; + let buffered_bytes = write_options.buffered_bytes_tracker(); let (bytes_written, handle) = match resolved { ResolvedStore::Path(path) => { @@ -466,7 +468,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create( arrow_schema, write_schema, bytes_written, - strategy, + buffered_bytes, handle, tx, )) @@ -513,10 +515,11 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_createStream( let strategy = write_strategy_for_schema(session, &write_schema); let write_options = session .write_options() - .with_strategy(Arc::clone(&strategy)) + .with_strategy(strategy) .with_metadata_segments(metadata); // See the note in `create`: validate before the write task can start. write_options.validate_metadata()?; + let buffered_bytes = write_options.buffered_bytes_tracker(); let mut write = CountingVortexWrite::new(JavaWrite::new(vm, writable)); let bytes_written = write.counter(); @@ -531,7 +534,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_createStream( arrow_schema, write_schema, bytes_written, - strategy, + buffered_bytes, handle, tx, )) diff --git a/vortex-layout/src/display.rs b/vortex-layout/src/display.rs index 20b13b7afbc..c3743a1df45 100644 --- a/vortex-layout/src/display.rs +++ b/vortex-layout/src/display.rs @@ -267,7 +267,7 @@ mod tests { let array1 = PrimitiveArray::new(buffer![1i64, 2, 3, 4, 5], validity); let layout1 = FlatLayoutStrategy::default() .write_stream( - ctx.clone(), + ctx.clone().into(), Arc::::clone(&segments), array1.into_array().to_array_stream().sequenced(ptr1), eof1, @@ -290,7 +290,7 @@ mod tests { } let layout2 = FlatLayoutStrategy::default() .write_stream( - ctx.clone(), + ctx.clone().into(), Arc::::clone(&segments), builder .finish() @@ -358,7 +358,7 @@ vortex.struct, dtype: {numbers=i64?, strings=utf8}, children: 2, rows: 5 PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5], Validity::NonNullable); let layout1 = FlatLayoutStrategy::default() .write_stream( - ctx.clone(), + ctx.clone().into(), Arc::::clone(&segments), array1.into_array().to_array_stream().sequenced(ptr1), eof1, @@ -373,7 +373,7 @@ vortex.struct, dtype: {numbers=i64?, strings=utf8}, children: 2, rows: 5 PrimitiveArray::new(buffer![6i32, 7, 8, 9, 10], Validity::NonNullable); let layout2 = FlatLayoutStrategy::default() .write_stream( - ctx.clone(), + ctx.clone().into(), Arc::::clone(&segments), array2.into_array().to_array_stream().sequenced(ptr2), eof2, @@ -421,7 +421,7 @@ vortex.chunked, dtype: i32, children: 2, rows: 10 let session = new_session().with_handle(handle); FlatLayoutStrategy::default() .write_stream( - ctx.clone(), + ctx.clone().into(), Arc::::clone(&segments), array.into_array().to_array_stream().sequenced(ptr), eof, @@ -466,7 +466,7 @@ vortex.flat, dtype: i32?, segment 0, buffers=[20B], total=20B let session = new_session().with_handle(handle); FlatLayoutStrategy::default() .write_stream( - ctx, + ctx.into(), Arc::::clone(&segments), array.into_array().to_array_stream().sequenced(ptr), eof, diff --git a/vortex-layout/src/layouts/buffered.rs b/vortex-layout/src/layouts/buffered.rs index aa5d64c67fd..b3a3e4652d9 100644 --- a/vortex-layout/src/layouts/buffered.rs +++ b/vortex-layout/src/layouts/buffered.rs @@ -3,19 +3,17 @@ use std::collections::VecDeque; use std::sync::Arc; -use std::sync::atomic::AtomicU64; -use std::sync::atomic::Ordering; use async_stream::try_stream; use async_trait::async_trait; use futures::StreamExt as _; use futures::pin_mut; -use vortex_array::ArrayContext; use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriterContext; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; use crate::sequence::SequencePointer; @@ -26,7 +24,6 @@ use crate::sequence::SequentialStreamExt as _; pub struct BufferedStrategy { child: Arc, buffer_size: u64, - buffered_bytes: Arc, } impl BufferedStrategy { @@ -34,7 +31,6 @@ impl BufferedStrategy { Self { child: Arc::new(child), buffer_size, - buffered_bytes: Arc::new(AtomicU64::new(0)), } } } @@ -43,7 +39,7 @@ impl BufferedStrategy { impl LayoutStrategy for BufferedStrategy { async fn write_stream( &self, - ctx: ArrayContext, + ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, stream: SendableSequentialStream, eof: SequencePointer, @@ -51,8 +47,8 @@ impl LayoutStrategy for BufferedStrategy { ) -> VortexResult { let dtype = stream.dtype().clone(); let buffer_size = self.buffer_size; + let buffered_bytes = ctx.buffered_bytes_tracker().clone(); - let buffered_bytes_counter = Arc::clone(&self.buffered_bytes); let buffered_stream = try_stream! { let stream = stream.peekable(); pin_mut!(stream); @@ -64,14 +60,13 @@ impl LayoutStrategy for BufferedStrategy { let (sequence_id, chunk) = chunk?; let chunk_size = chunk.nbytes(); nbytes += chunk_size; - buffered_bytes_counter.fetch_add(chunk_size, Ordering::Relaxed); - chunks.push_back(chunk); + chunks.push_back((chunk, buffered_bytes.reserve(chunk_size))); // 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) = chunks.pop_front() { - buffered_bytes_counter.fetch_sub(chunk.nbytes(), Ordering::Relaxed); + while let Some((chunk, reservation)) = chunks.pop_front() { + drop(reservation); yield (sequence_ptr.advance(), chunk) } break; @@ -85,12 +80,11 @@ impl LayoutStrategy for BufferedStrategy { // 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) = chunks.pop_front() else { + let Some((chunk, reservation)) = chunks.pop_front() else { break; }; - let chunk_size = chunk.nbytes(); - nbytes -= chunk_size; - buffered_bytes_counter.fetch_sub(chunk_size, Ordering::Relaxed); + nbytes -= reservation.bytes(); + drop(reservation); yield (sequence_ptr.advance(), chunk) } } @@ -106,8 +100,4 @@ impl LayoutStrategy for BufferedStrategy { ) .await } - - fn buffered_bytes(&self) -> u64 { - self.buffered_bytes.load(Ordering::Relaxed) + self.child.buffered_bytes() - } } diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs index 0e33f39df0c..978dad4cbe8 100644 --- a/vortex-layout/src/layouts/chunked/reader.rs +++ b/vortex-layout/src/layouts/chunked/reader.rs @@ -395,7 +395,7 @@ mod test { let session = new_session().with_handle(handle); strategy .write_stream( - ctx, + ctx.into(), segments2, SequentialStreamAdapter::new( DType::Primitive(PType::I32, NonNullable), diff --git a/vortex-layout/src/layouts/chunked/writer.rs b/vortex-layout/src/layouts/chunked/writer.rs index 8cc4536588d..8de7c30fe47 100644 --- a/vortex-layout/src/layouts/chunked/writer.rs +++ b/vortex-layout/src/layouts/chunked/writer.rs @@ -8,7 +8,6 @@ use async_trait::async_trait; use futures::StreamExt; use futures::TryStreamExt; use futures::stream; -use vortex_array::ArrayContext; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_io::session::RuntimeSessionExt; @@ -16,6 +15,7 @@ use vortex_session::VortexSession; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriterContext; use crate::children::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; use crate::segments::SegmentSinkRef; @@ -42,7 +42,7 @@ impl ChunkedLayoutStrategy { impl LayoutStrategy for ChunkedLayoutStrategy { async fn write_stream( &self, - ctx: ArrayContext, + ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, stream: SendableSequentialStream, mut eof: SequencePointer, @@ -99,8 +99,4 @@ impl LayoutStrategy for ChunkedLayoutStrategy { .into_layout()) } } - - fn buffered_bytes(&self) -> u64 { - self.chunk_strategy.buffered_bytes() - } } diff --git a/vortex-layout/src/layouts/collect.rs b/vortex-layout/src/layouts/collect.rs index 21dd25c756c..1fad03ef0a7 100644 --- a/vortex-layout/src/layouts/collect.rs +++ b/vortex-layout/src/layouts/collect.rs @@ -7,7 +7,6 @@ use async_stream::try_stream; use async_trait::async_trait; use futures::StreamExt; use futures::pin_mut; -use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; use vortex_error::VortexResult; @@ -15,6 +14,7 @@ use vortex_session::VortexSession; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriterContext; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; use crate::sequence::SequencePointer; @@ -39,7 +39,7 @@ impl CollectStrategy { impl LayoutStrategy for CollectStrategy { async fn write_stream( &self, - ctx: ArrayContext, + ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, stream: SendableSequentialStream, eof: SequencePointer, @@ -75,8 +75,4 @@ impl LayoutStrategy for CollectStrategy { .write_stream(ctx, segment_sink, adapted, eof, session) .await } - - fn buffered_bytes(&self) -> u64 { - todo!() - } } diff --git a/vortex-layout/src/layouts/compressed.rs b/vortex-layout/src/layouts/compressed.rs index efd85db9e5d..87af23bb469 100644 --- a/vortex-layout/src/layouts/compressed.rs +++ b/vortex-layout/src/layouts/compressed.rs @@ -5,7 +5,6 @@ use std::sync::Arc; use async_trait::async_trait; use futures::StreamExt as _; -use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::VortexSessionExecute; @@ -18,6 +17,7 @@ use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriterContext; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; use crate::sequence::SequencePointer; @@ -89,7 +89,7 @@ impl CompressingStrategy { impl LayoutStrategy for CompressingStrategy { async fn write_stream( &self, - ctx: ArrayContext, + ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, stream: SendableSequentialStream, eof: SequencePointer, @@ -127,8 +127,4 @@ impl LayoutStrategy for CompressingStrategy { ) .await } - - fn buffered_bytes(&self) -> u64 { - self.child.buffered_bytes() - } } diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index d0a96d56f4b..36d4bf04abf 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -415,7 +415,7 @@ mod tests { let dtype = array.dtype().clone(); let layout = strategy .write_stream( - ArrayContext::empty(), + ArrayContext::empty().into(), Arc::::clone(&segments), SequentialStreamAdapter::new(dtype, array.to_array_stream().sequenced(ptr)) .sendable(), @@ -462,7 +462,7 @@ mod tests { let (ptr, eof) = SequenceId::root().split(); let layout: LayoutRef = strategy .write_stream( - array_ctx, + array_ctx.into(), Arc::::clone(&segments), SequentialStreamAdapter::new( DType::Utf8(Nullability::Nullable), @@ -549,7 +549,7 @@ mod tests { let (ptr, eof) = SequenceId::root().split(); let layout: LayoutRef = strategy .write_stream( - array_ctx, + array_ctx.into(), Arc::::clone(&segments), SequentialStreamAdapter::new( DType::Utf8(Nullability::Nullable), @@ -617,7 +617,7 @@ mod tests { let (ptr, eof) = SequenceId::root().split(); let layout: LayoutRef = strategy .write_stream( - array_ctx, + array_ctx.into(), Arc::::clone(&segments), SequentialStreamAdapter::new( DType::Utf8(Nullability::Nullable), diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index c93142d2f15..b7c93a992fb 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -18,7 +18,6 @@ use futures::pin_mut; use futures::stream::BoxStream; use futures::stream::once; use futures::try_join; -use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -40,6 +39,7 @@ use vortex_session::VortexSession; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriterContext; use crate::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; use crate::layouts::compressed::CompressorPlugin; @@ -132,7 +132,7 @@ impl DictStrategy { impl LayoutStrategy for DictStrategy { async fn write_stream( &self, - ctx: ArrayContext, + ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, stream: SendableSequentialStream, mut eof: SequencePointer, @@ -251,10 +251,6 @@ impl LayoutStrategy for DictStrategy { ) .into_layout()) } - - fn buffered_bytes(&self) -> u64 { - self.codes.buffered_bytes() + self.values.buffered_bytes() + self.fallback.buffered_bytes() - } } enum DictionaryChunk { diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index 42e707155a8..12191b35429 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -265,7 +265,7 @@ mod test { PrimitiveArray::new(buffer![1, 2, 3, 4, 5], Validity::AllValid).into_array(); let layout = FlatLayoutStrategy::default() .write_stream( - array_ctx, + array_ctx.into(), Arc::::clone(&segments), array.to_array_stream().sequenced(ptr), eof, @@ -306,7 +306,7 @@ mod test { PrimitiveArray::new(buffer![1, 2, 3, 4, 5], Validity::AllValid).into_array(); let layout = FlatLayoutStrategy::default() .write_stream( - array_ctx, + array_ctx.into(), Arc::::clone(&segments), array.to_array_stream().sequenced(ptr), eof, @@ -345,7 +345,7 @@ mod test { PrimitiveArray::new(buffer![1, 2, 3, 4, 5], Validity::AllValid).into_array(); let layout = FlatLayoutStrategy::default() .write_stream( - array_ctx, + array_ctx.into(), Arc::::clone(&segments), array.to_array_stream().sequenced(ptr), eof, diff --git a/vortex-layout/src/layouts/flat/writer.rs b/vortex-layout/src/layouts/flat/writer.rs index efb2c11cbe3..83c2101f12c 100644 --- a/vortex-layout/src/layouts/flat/writer.rs +++ b/vortex-layout/src/layouts/flat/writer.rs @@ -3,7 +3,6 @@ use async_trait::async_trait; use futures::StreamExt; -use vortex_array::ArrayContext; use vortex_array::dtype::DType; use vortex_array::expr::stats::Precision; use vortex_array::expr::stats::Stat; @@ -24,6 +23,7 @@ use vortex_session::registry::ReadContext; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriterContext; use crate::children::OwnedLayoutChildren; use crate::layouts::chunked::ChunkedLayout; use crate::layouts::flat::FlatLayout; @@ -83,13 +83,12 @@ fn truncate_scalar_stat Option<(Scalar, bool)>>( impl LayoutStrategy for FlatLayoutStrategy { async fn write_stream( &self, - ctx: ArrayContext, + ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, mut stream: SendableSequentialStream, _eof: SequencePointer, session: &VortexSession, ) -> VortexResult { - let ctx = ctx.clone(); let Some(chunk) = stream.next().await else { // an empty input has no segment to write. return Ok(ChunkedLayout::new( @@ -144,7 +143,7 @@ impl LayoutStrategy for FlatLayoutStrategy { } let buffers = chunk.serialize( - &ctx, + ctx.array_ctx(), session, &SerializeOptions { offset: 0, @@ -164,16 +163,11 @@ impl LayoutStrategy for FlatLayoutStrategy { row_count, stream.dtype().clone(), segment_id, - ReadContext::new(ctx.to_ids()), + ReadContext::new(ctx.array_ctx().to_ids()), array_node, ) .into_layout()) } - - fn buffered_bytes(&self) -> u64 { - // FlatLayoutStrategy is a leaf strategy with no child strategies and no buffering - 0 - } } #[cfg(test)] @@ -236,7 +230,7 @@ mod tests { let array = PrimitiveArray::new(buffer![1, 2, 3, 4, 5], Validity::AllValid); let layout = FlatLayoutStrategy::default() .write_stream( - ctx, + ctx.into(), Arc::::clone(&segments), array.into_array().to_array_stream().sequenced(ptr), eof, @@ -287,7 +281,7 @@ mod tests { let layout = FlatLayoutStrategy::default() .write_stream( - ctx, + ctx.into(), Arc::::clone(&segments), array.into_array().to_array_stream().sequenced(ptr), eof, @@ -356,7 +350,7 @@ mod tests { let (ptr, eof) = SequenceId::root().split(); let layout = FlatLayoutStrategy::default() .write_stream( - ctx, + ctx.into(), Arc::::clone(&segments), array.into_array().to_array_stream().sequenced(ptr), eof, @@ -430,7 +424,7 @@ mod tests { let layout = LayoutStrategyEncodingValidator::new(FlatLayoutStrategy::default(), allowed) .write_stream( - ctx, + ctx.into(), Arc::::clone(&segments), filter.into_array().to_array_stream().sequenced(ptr), eof, @@ -472,7 +466,7 @@ mod tests { let layout = LayoutStrategyEncodingValidator::new(FlatLayoutStrategy::default(), allowed) .write_stream( - ctx, + ctx.into(), Arc::::clone(&segments), dict.into_array().to_array_stream().sequenced(ptr), eof, diff --git a/vortex-layout/src/layouts/list/reader.rs b/vortex-layout/src/layouts/list/reader.rs index 1e21714fdd0..4bea8e74ca0 100644 --- a/vortex-layout/src/layouts/list/reader.rs +++ b/vortex-layout/src/layouts/list/reader.rs @@ -841,7 +841,13 @@ mod tests { let (ptr, eof) = SequenceId::root().split(); let stream = array.to_array_stream().sequenced(ptr); let layout = strategy - .write_stream(ArrayContext::empty(), segments, stream, eof, &session) + .write_stream( + ArrayContext::empty().into(), + segments, + stream, + eof, + &session, + ) .await?; Ok((segments_ref, layout, session)) } diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index f9f0516df23..4d8565fdd10 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -7,7 +7,6 @@ use async_trait::async_trait; use futures::StreamExt; use futures::future::try_join; use futures::future::try_join_all; -use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -33,6 +32,7 @@ use vortex_session::VortexSession; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriterContext; use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::layouts::list::ListLayout; use crate::segments::SegmentSinkRef; @@ -115,7 +115,7 @@ impl ListLayoutStrategy { impl LayoutStrategy for ListLayoutStrategy { async fn write_stream( &self, - ctx: ArrayContext, + ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, stream: SendableSequentialStream, mut eof: SequencePointer, @@ -205,13 +205,6 @@ impl LayoutStrategy for ListLayoutStrategy { Ok(ListLayout::new(dtype, elements_layout, offsets_layout, validity_layout).into_layout()) } - - fn buffered_bytes(&self) -> u64 { - let list_bytes = self.elements.buffered_bytes() - + self.offsets.buffered_bytes() - + self.validity.buffered_bytes(); - list_bytes.max(self.fallback.buffered_bytes()) - } } /// Transpose a list column into its `elements`, `offsets`, and (when present) `validity` child @@ -330,6 +323,7 @@ impl Matcher for AnyList { #[cfg(test)] mod tests { use futures::stream; + use vortex_array::ArrayContext; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ListArray; @@ -365,7 +359,13 @@ mod tests { let (ptr, eof) = SequenceId::root().split(); let stream = array.to_array_stream().sequenced(ptr); strategy - .write_stream(ArrayContext::empty(), segments, stream, eof, &session) + .write_stream( + ArrayContext::empty().into(), + segments, + stream, + eof, + &session, + ) .await } @@ -441,7 +441,13 @@ mod tests { let session = layout_test_session(); let res = flat_list_strategy() - .write_stream(ArrayContext::empty(), segments, stream, eof, &session) + .write_stream( + ArrayContext::empty().into(), + segments, + stream, + eof, + &session, + ) .await; assert!(res.is_err()) } diff --git a/vortex-layout/src/layouts/repartition.rs b/vortex-layout/src/layouts/repartition.rs index c0da648a76a..7344ee7be51 100644 --- a/vortex-layout/src/layouts/repartition.rs +++ b/vortex-layout/src/layouts/repartition.rs @@ -8,7 +8,6 @@ use async_stream::try_stream; use async_trait::async_trait; use futures::StreamExt as _; use futures::pin_mut; -use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; @@ -21,6 +20,7 @@ use vortex_session::VortexSession; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriterContext; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; use crate::sequence::SequencePointer; @@ -95,7 +95,7 @@ impl RepartitionStrategy { impl LayoutStrategy for RepartitionStrategy { async fn write_stream( &self, - ctx: ArrayContext, + ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, stream: SendableSequentialStream, eof: SequencePointer, @@ -186,14 +186,6 @@ impl LayoutStrategy for RepartitionStrategy { ) .await } - - fn buffered_bytes(&self) -> u64 { - // TODO(os): we should probably add the buffered bytes from this strategy on top, - // it is currently better to not add it at all because these buffered arrays are - // potentially sliced and uncompressed. They would overestimate the actual bytes - // that will end up in the file when flushed. - self.child.buffered_bytes() - } } struct ChunksBuffer { @@ -401,7 +393,7 @@ mod tests { let session = new_session().with_handle(handle); strategy .write_stream( - ctx, + ctx.into(), Arc::::clone(&segments), stream, eof, @@ -468,7 +460,7 @@ mod tests { let session = new_session().with_handle(handle); strategy .write_stream( - ctx, + ctx.into(), Arc::::clone(&segments), stream, eof, diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index fe94d1b9e77..7eaa820ce2a 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -360,7 +360,7 @@ mod tests { let array = buffer![1..=5].into_array(); let layout = FlatLayoutStrategy::default() .write_stream( - array_ctx, + array_ctx.into(), Arc::::clone(&segments), array.to_array_stream().sequenced(ptr), eof, @@ -405,7 +405,7 @@ mod tests { let array = buffer![1..=5].into_array(); let layout = FlatLayoutStrategy::default() .write_stream( - array_ctx, + array_ctx.into(), Arc::::clone(&segments), array.to_array_stream().sequenced(ptr), eof, @@ -450,7 +450,7 @@ mod tests { let array = buffer![1..=5].into_array(); let layout = FlatLayoutStrategy::default() .write_stream( - array_ctx, + array_ctx.into(), Arc::::clone(&segments), array.to_array_stream().sequenced(ptr), eof, diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 7b7754fa42b..ee939ecd6bd 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -460,7 +460,7 @@ mod tests { let session = new_session().with_handle(handle); strategy .write_stream( - ctx, + ctx.into(), segments2, StructArray::try_new( Vec::::new().into(), @@ -497,7 +497,7 @@ mod tests { let session = new_session().with_handle(handle); strategy .write_stream( - ctx, + ctx.into(), segments2, StructArray::from_fields( [ @@ -537,7 +537,7 @@ mod tests { let session = new_session().with_handle(handle); strategy .write_stream( - ctx, + ctx.into(), segments2, StructArray::try_from_iter_with_validity( [ @@ -582,7 +582,7 @@ mod tests { let session = new_session().with_handle(handle); strategy .write_stream( - ctx, + ctx.into(), segments2, StructArray::try_from_iter_with_validity( [( @@ -862,7 +862,7 @@ mod tests { let session = new_session().with_handle(handle); strategy .write_stream( - ctx, + ctx.into(), segments2, StructArray::from_fields( [ diff --git a/vortex-layout/src/layouts/struct_/writer.rs b/vortex-layout/src/layouts/struct_/writer.rs index cc2e239144f..8db5f2b8205 100644 --- a/vortex-layout/src/layouts/struct_/writer.rs +++ b/vortex-layout/src/layouts/struct_/writer.rs @@ -20,7 +20,6 @@ use futures::future::try_join; use futures::future::try_join_all; use futures::pin_mut; use itertools::Itertools; -use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -41,6 +40,7 @@ use vortex_utils::aliases::hash_set::HashSet; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriterContext; use crate::layouts::struct_::StructLayout; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; @@ -105,7 +105,7 @@ impl StructStrategy { impl LayoutStrategy for StructStrategy { async fn write_stream( &self, - ctx: ArrayContext, + ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, stream: SendableSequentialStream, mut eof: SequencePointer, diff --git a/vortex-layout/src/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 8fade569ad5..1a3c1adc524 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -17,7 +17,6 @@ use std::sync::Arc; use std::sync::LazyLock; use async_trait::async_trait; -use vortex_array::ArrayContext; use vortex_array::dtype::Field; use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldPath; @@ -28,6 +27,7 @@ use vortex_utils::aliases::hash_set::HashSet; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriterContext; use crate::layouts::list::writer::ListLayoutStrategy; use crate::layouts::struct_::StructStrategy; use crate::segments::SegmentSinkRef; @@ -38,7 +38,7 @@ use crate::sequence::SequencePointer; /// default. Disabled unless the environment variable `VORTEX_EXPERIMENTAL_LIST_LAYOUT` /// is set to `1`. /// -/// [`ListLayoutStrategy`]: crate::layouts::list::writer::ListLayoutStrategy +/// [`ListLayoutStrategy`]: ListLayoutStrategy pub fn use_experimental_list_layout() -> bool { static USE_EXPERIMENTAL_LIST_LAYOUT: LazyLock = LazyLock::new(|| env::var("VORTEX_EXPERIMENTAL_LIST_LAYOUT").is_ok_and(|v| v == "1")); @@ -72,7 +72,7 @@ pub struct TableStrategy { /// Optional factory applied to each dynamically constructed [`ListLayoutStrategy`]. /// Its presence also enables list decomposition. /// - /// [`ListLayoutStrategy`]: crate::layouts::list::writer::ListLayoutStrategy + /// [`ListLayoutStrategy`]: ListLayoutStrategy list_layout_factory: Option, } @@ -295,7 +295,7 @@ impl TableStrategy { impl LayoutStrategy for TableStrategy { async fn write_stream( &self, - ctx: ArrayContext, + ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, stream: SendableSequentialStream, eof: SequencePointer, @@ -376,7 +376,13 @@ mod tests { let stream = array.to_array_stream().sequenced(ptr); let session = new_session().with_tokio(); strategy - .write_stream(ArrayContext::empty(), segments, stream, eof, &session) + .write_stream( + ArrayContext::empty().into(), + segments, + stream, + eof, + &session, + ) .await } @@ -685,7 +691,7 @@ mod tests { strategy .write_stream( - ctx, + ctx.into(), segments, SequentialStreamAdapter::new(dtype, stream).sendable(), eof, diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index f2b83836352..dbcef1f7a68 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -311,7 +311,7 @@ mod test { let layout = block_on(|handle| async move { let session = session_with_handle(handle); strategy - .write_stream(ctx, segments2, array_stream, eof, &session) + .write_stream(ctx.into(), segments2, array_stream, eof, &session) .await }) .unwrap(); @@ -408,7 +408,7 @@ mod test { let layout = block_on(|handle| async move { let session = session_with_handle(handle); strategy - .write_stream(ctx, segments2, array_stream, eof, &session) + .write_stream(ctx.into(), segments2, array_stream, eof, &session) .await }) .unwrap(); diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index 77d7c770319..341eded62fb 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -9,7 +9,6 @@ use std::sync::Arc; use async_trait::async_trait; use futures::StreamExt as _; use parking_lot::Mutex; -use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; @@ -36,6 +35,7 @@ use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutRef; use crate::LayoutStrategy; +use crate::LayoutWriterContext; use crate::layouts::zoned::AggregateStatsAccumulator; use crate::layouts::zoned::ZonedLayout; use crate::layouts::zoned::aggregate_partials; @@ -99,7 +99,7 @@ impl ZonedStrategy { impl LayoutStrategy for ZonedStrategy { async fn write_stream( &self, - ctx: ArrayContext, + ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, stream: SendableSequentialStream, mut eof: SequencePointer, @@ -189,10 +189,6 @@ impl LayoutStrategy for ZonedStrategy { .into_layout(), ) } - - fn buffered_bytes(&self) -> u64 { - self.child.buffered_bytes() + self.stats.buffered_bytes() - } } fn default_zoned_aggregate_fns(dtype: &DType, session: &VortexSession) -> Arc<[AggregateFnRef]> { diff --git a/vortex-layout/src/scan/split_by.rs b/vortex-layout/src/scan/split_by.rs index 6c2cde11b72..be145572524 100644 --- a/vortex-layout/src/scan/split_by.rs +++ b/vortex-layout/src/scan/split_by.rs @@ -157,7 +157,7 @@ mod test { let session = session_with_handle(handle); FlatLayoutStrategy::default() .write_stream( - ctx, + ctx.into(), Arc::::clone(&segments), buffer![1_i32; 10] .into_array() diff --git a/vortex-layout/src/strategy.rs b/vortex-layout/src/strategy.rs index f9f91d9d82f..7a582418fee 100644 --- a/vortex-layout/src/strategy.rs +++ b/vortex-layout/src/strategy.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; use async_trait::async_trait; use futures::StreamExt; @@ -20,7 +22,109 @@ use crate::sequence::SequencePointer; use crate::sequence::SequentialStreamAdapter; use crate::sequence::SequentialStreamExt; -// [layout writer] +/// A shared counter of the bytes that layout strategies are holding but have not yet emitted. +/// +/// 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. +#[derive(Clone, Debug, Default)] +pub struct BufferedBytesTracker(Arc); + +impl BufferedBytesTracker { + /// Creates a tracker with a zeroed counter. + pub fn new() -> Self { + Self::default() + } + + /// Returns the number of bytes currently retained by layout strategies. + pub fn buffered_bytes(&self) -> u64 { + self.0.load(Ordering::Relaxed) + } + + /// Records `bytes` as buffered until the returned reservation is dropped. + pub fn reserve(&self, bytes: u64) -> BufferedBytesReservation { + self.0.fetch_add(bytes, Ordering::Relaxed); + BufferedBytesReservation { + tracker: self.clone(), + bytes, + } + } +} + +/// An outstanding claim on a [`BufferedBytesTracker`], released when dropped. +#[derive(Debug)] +pub struct BufferedBytesReservation { + tracker: BufferedBytesTracker, + bytes: u64, +} + +impl BufferedBytesReservation { + /// Returns the number of bytes held by this reservation. + pub fn bytes(&self) -> u64 { + self.bytes + } +} + +impl Drop for BufferedBytesReservation { + fn drop(&mut self) { + self.tracker.0.fetch_sub(self.bytes, Ordering::Relaxed); + } +} + +/// State shared by every strategy participating in a single layout write. +/// +/// Clones share the [`BufferedBytesTracker`] while retaining the array serialization context. +/// Passing this context through the strategy tree keeps writer-scoped state independent of the +/// strategy instances, which may be shared by multiple leaves or writers. +#[derive(Clone)] +pub struct LayoutWriterContext { + array_ctx: ArrayContext, + buffered_bytes: BufferedBytesTracker, +} + +impl LayoutWriterContext { + /// Creates a context for a layout write with a fresh buffered bytes tracker. + pub fn new(array_ctx: ArrayContext) -> Self { + Self { + array_ctx, + buffered_bytes: BufferedBytesTracker::new(), + } + } + + /// Replaces the buffered bytes tracker, so callers can observe the counter from outside the + /// strategy tree. + pub fn with_buffered_bytes_tracker(mut self, tracker: BufferedBytesTracker) -> Self { + self.buffered_bytes = tracker; + self + } + + /// Returns the array serialization context. + pub fn array_ctx(&self) -> &ArrayContext { + &self.array_ctx + } + + /// Returns the tracker that accounts for bytes retained by layout strategies. + pub fn buffered_bytes_tracker(&self) -> &BufferedBytesTracker { + &self.buffered_bytes + } + + /// Returns the number of bytes currently retained by layout strategies. + pub fn buffered_bytes(&self) -> u64 { + self.buffered_bytes.buffered_bytes() + } + + /// Records `bytes` as retained by this write until the returned reservation is dropped. + pub fn reserve_buffered_bytes(&self, bytes: u64) -> BufferedBytesReservation { + self.buffered_bytes.reserve(bytes) + } +} + +impl From for LayoutWriterContext { + fn from(array_ctx: ArrayContext) -> Self { + Self::new(array_ctx) + } +} + /// Writes an ordered array stream into a layout tree and segment sink. /// /// Layout strategies are writer-side extension points. Strategies may repartition, buffer, @@ -44,6 +148,8 @@ pub trait LayoutStrategy: 'static + Send + Sync { /// 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. /// @@ -63,21 +169,12 @@ pub trait LayoutStrategy: 'static + Send + Sync { /// of data, or serializing very large messages to flatbuffers. async fn write_stream( &self, - ctx: ArrayContext, + ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, stream: SendableSequentialStream, eof: SequencePointer, session: &VortexSession, ) -> VortexResult; - - /// Returns the number of bytes currently buffered by this strategy and any child strategies. - /// - /// This method allows tracking of data that has been processed by the strategy but not yet - /// written to the underlying sink, providing more accurate estimates of final file size - /// during write operations. - fn buffered_bytes(&self) -> u64 { - 0 - } } /// A layout strategy wrapper that rejects arrays containing encodings outside an allow-list. @@ -104,7 +201,7 @@ impl LayoutStrategyEncodingValidator { impl LayoutStrategy for LayoutStrategyEncodingValidator { async fn write_stream( &self, - ctx: ArrayContext, + ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, stream: SendableSequentialStream, eof: SequencePointer, @@ -131,17 +228,13 @@ impl LayoutStrategy for LayoutStrategyEncodingValidator { ) .await } - - fn buffered_bytes(&self) -> u64 { - self.child.buffered_bytes() - } } #[async_trait] impl LayoutStrategy for Arc { async fn write_stream( &self, - ctx: ArrayContext, + ctx: LayoutWriterContext, segment_sink: SegmentSinkRef, stream: SendableSequentialStream, eof: SequencePointer, @@ -151,9 +244,38 @@ impl LayoutStrategy for Arc { .write_stream(ctx, segment_sink, stream, eof, session) .await } +} + +#[cfg(test)] +mod tests { + use crate::strategy::BufferedBytesTracker; + + #[test] + fn reservations_accumulate_and_release() { + let tracker = BufferedBytesTracker::new(); + assert_eq!(tracker.buffered_bytes(), 0); + + let first = tracker.reserve(16); + let second = tracker.reserve(32); + assert_eq!(tracker.buffered_bytes(), 48); + assert_eq!(first.bytes(), 16); + + drop(first); + assert_eq!(tracker.buffered_bytes(), 32); + + drop(second); + assert_eq!(tracker.buffered_bytes(), 0); + } + + #[test] + fn clones_share_the_same_counter() { + let tracker = BufferedBytesTracker::new(); + let observer = tracker.clone(); + + let reservation = tracker.reserve(8); + assert_eq!(observer.buffered_bytes(), 8); - fn buffered_bytes(&self) -> u64 { - (**self).buffered_bytes() + drop(reservation); + assert_eq!(observer.buffered_bytes(), 0); } } -// [layout writer]