Replace LayoutStrategy buffered_bytes with a buffered bytes tracker

vortex-array/src/serde.rs

@@ -191,16 +191,12 @@ impl<'a> ArrayNodeFlatBuffer<'a> {
         &self,
         fbb: &mut FlatBufferBuilder<'fb>,
     ) -> VortexResult<WIPOffset<fba::ArrayNode<'fb>>> {
-        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!(

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<F: Fn(Scalar) -> 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<LayoutRef> {
-        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),
             },

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<dyn LayoutStrategy> {
+        Arc::new(BufferedStrategy::new(
+            ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()),
+            BUFFER_SIZE,
+        ))
+    };
+    let validity: Arc<dyn LayoutStrategy> = Arc::new(FlatLayoutStrategy::default());
+    let strategy: Arc<dyn LayoutStrategy> = 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<dyn LayoutStrategy> = 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::<EditionSession>()
+        .with::<LayoutSession>()
+        .with::<RuntimeSession>();
+
+    // 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

vortex-file/src/v2/file_stats_reader.rs

@@ -246,7 +246,7 @@ mod tests {
             );
             let layout = strategy
                 .write_stream(
-                    ctx,
+                    ctx.into(),
                     Arc::<TestSegments>::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::<TestSegments>::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::<TestSegments>::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<dyn SegmentSink>,
                     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<dyn SegmentSink>,
                     struct_array.into_array().to_array_stream().sequenced(ptr),
                     eof,