buffered strategy to not use eof for the final chunk (#7219)

## Summary

fixes buffered layout writer so it doesn't write the final chunk on the
eof pointer. Eof should only be used for data that the writer wants to
place at the end of the file. Buffered writer was writing regular
buffered data to there which did mess up ordering of some segments.

Previously struct writer was using a transposed stream without spawning
a task per column, on that world buffering was deadlocky. That is
changed for a while to spawn now, so we should be deadlock safe.

I did try converting all clickbench files repeatedly, as well as the
public bi datasets and randomly generated wide tables but I couldn't
deadlock this.

fixes #7234
fixes #7236

## Testing

add a vortex file test that asserts the dict layout segments are in the
right order, as well as zone maps across columns

---------

Signed-off-by: Onur Satici <onur@spiraldb.com>
Signed-off-by: Will Manning <will@willmanning.io>

Author: onursatici

vortex-file/src/tests.rs

@@ -62,6 +62,7 @@ use vortex_buffer::ByteBufferMut;
 use vortex_buffer::buffer;
 use vortex_error::VortexResult;
 use vortex_io::session::RuntimeSession;
+use vortex_layout::Layout;
 use vortex_layout::scan::scan_builder::ScanBuilder;
 use vortex_layout::session::LayoutSession;
 use vortex_session::VortexSession;
@@ -71,6 +72,7 @@ use crate::V1_FOOTER_FBS_SIZE;
 use crate::VERSION;
 use crate::VortexFile;
 use crate::WriteOptionsSessionExt;
+use crate::footer::SegmentSpec;
 
 static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
     let mut session = VortexSession::empty()
@@ -1696,3 +1698,176 @@ async fn timestamp_unit_mismatch_errors_with_constant_children()
 
     Ok(())
 }
+
+/// Collect all segment byte offsets reachable from a layout node.
+fn collect_segment_offsets(layout: &dyn Layout, segment_specs: &[SegmentSpec]) -> Vec<u64> {
+    let mut result = Vec::new();
+    collect_segment_offsets_inner(layout, segment_specs, &mut result);
+    result
+}
+
+fn collect_segment_offsets_inner(
+    layout: &dyn Layout,
+    segment_specs: &[SegmentSpec],
+    result: &mut Vec<u64>,
+) {
+    for seg_id in layout.segment_ids() {
+        result.push(segment_specs[*seg_id as usize].offset);
+    }
+    for child in layout.children().unwrap() {
+        collect_segment_offsets_inner(child.as_ref(), segment_specs, result);
+    }
+}
+
+/// Assert that all offsets in `before` are less than all offsets in `after`.
+fn assert_offsets_ordered(before: &[u64], after: &[u64], context: &str) {
+    if let (Some(&max_before), Some(&min_after)) = (before.iter().max(), after.iter().min()) {
+        assert!(
+            max_before < min_after,
+            "{context}: expected all 'before' offsets < all 'after' offsets, \
+             but max before = {max_before} >= min after = {min_after}"
+        );
+    }
+}
+
+#[tokio::test]
+#[cfg_attr(miri, ignore)]
+async fn test_segment_ordering_dict_codes_before_values() -> VortexResult<()> {
+    // Create low-cardinality strings to trigger dict encoding, plus an integer column.
+    let n = 100_000;
+    let values: Vec<&str> = (0..n).map(|i| ["alpha", "beta", "gamma"][i % 3]).collect();
+    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 mut buf = ByteBufferMut::empty();
+    let summary = SESSION
+        .write_options()
+        .write(&mut buf, st.to_array_stream())
+        .await?;
+
+    let footer = summary.footer();
+    let segment_specs = footer.segment_map();
+    let root = footer.layout();
+
+    // Walk the layout tree and find all dict layouts.
+    // Verify codes segments come before values segments in byte order within each run.
+    fn check_dict_ordering(layout: &dyn Layout, segment_specs: &[SegmentSpec]) {
+        if layout.encoding_id().as_ref() == "vortex.dict" {
+            // child 0 = values, child 1 = codes
+            let values_offsets =
+                collect_segment_offsets(layout.child(0).unwrap().as_ref(), segment_specs);
+            let codes_offsets =
+                collect_segment_offsets(layout.child(1).unwrap().as_ref(), segment_specs);
+
+            assert_offsets_ordered(
+                &codes_offsets,
+                &values_offsets,
+                "dict: codes should come before values",
+            );
+        }
+
+        for child in layout.children().unwrap() {
+            check_dict_ordering(child.as_ref(), segment_specs);
+        }
+    }
+
+    check_dict_ordering(root.as_ref(), segment_specs);
+
+    Ok(())
+}
+
+#[tokio::test]
+#[cfg_attr(miri, ignore)]
+async fn test_segment_ordering_zonemaps_after_data() -> VortexResult<()> {
+    // Create a multi-column struct with enough rows to produce zone maps.
+    let n = 100_000;
+    let values: Vec<&str> = (0..n).map(|i| ["alpha", "beta", "gamma"][i % 3]).collect();
+    let strings = VarBinArray::from(values).into_array();
+    let numbers = PrimitiveArray::from_iter(0..n as i32).into_array();
+    let floats = PrimitiveArray::from_iter((0..n).map(|i| i as f64 * 0.1)).into_array();
+
+    let st = StructArray::from_fields(&[
+        ("strings", strings),
+        ("numbers", numbers),
+        ("floats", floats),
+    ])
+    .unwrap();
+
+    let mut buf = ByteBufferMut::empty();
+    let summary = SESSION
+        .write_options()
+        .write(&mut buf, st.to_array_stream())
+        .await?;
+
+    let footer = summary.footer();
+    let segment_specs = footer.segment_map();
+    let root = footer.layout();
+
+    // Find all zoned layouts and verify data segments come before zone map segments.
+    fn check_zoned_ordering(layout: &dyn Layout, segment_specs: &[SegmentSpec]) {
+        if layout.encoding_id().as_ref() == "vortex.stats" {
+            // child 0 = data, child 1 = zones
+            let data_offsets =
+                collect_segment_offsets(layout.child(0).unwrap().as_ref(), segment_specs);
+            let zones_offsets =
+                collect_segment_offsets(layout.child(1).unwrap().as_ref(), segment_specs);
+
+            assert_offsets_ordered(
+                &data_offsets,
+                &zones_offsets,
+                "zoned: data should come before zones",
+            );
+        }
+
+        for child in layout.children().unwrap() {
+            check_zoned_ordering(child.as_ref(), segment_specs);
+        }
+    }
+
+    check_zoned_ordering(root.as_ref(), segment_specs);
+
+    // Additionally: all zone map segments across all columns should appear after
+    // all data segments across all columns.
+    let mut all_data_offsets = Vec::new();
+    let mut all_zones_offsets = Vec::new();
+
+    fn collect_all_zoned(
+        layout: &dyn Layout,
+        segment_specs: &[SegmentSpec],
+        all_data: &mut Vec<u64>,
+        all_zones: &mut Vec<u64>,
+    ) {
+        if layout.encoding_id().as_ref() == "vortex.stats" {
+            // child 0 = data, child 1 = zones
+            all_data.extend(collect_segment_offsets(
+                layout.child(0).unwrap().as_ref(),
+                segment_specs,
+            ));
+            all_zones.extend(collect_segment_offsets(
+                layout.child(1).unwrap().as_ref(),
+                segment_specs,
+            ));
+            return;
+        }
+        for child in layout.children().unwrap() {
+            collect_all_zoned(child.as_ref(), segment_specs, all_data, all_zones);
+        }
+    }
+
+    collect_all_zoned(
+        root.as_ref(),
+        segment_specs,
+        &mut all_data_offsets,
+        &mut all_zones_offsets,
+    );
+
+    assert_offsets_ordered(
+        &all_data_offsets,
+        &all_zones_offsets,
+        "global: all data segments should come before all zone map segments",
+    );
+
+    Ok(())
+}

vortex-layout/src/layouts/buffered.rs

@@ -9,6 +9,7 @@ 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_io::runtime::Handle;
@@ -44,20 +45,18 @@ impl LayoutStrategy for BufferedStrategy {
         &self,
         ctx: ArrayContext,
         segment_sink: SegmentSinkRef,
-        mut stream: SendableSequentialStream,
-        mut eof: SequencePointer,
+        stream: SendableSequentialStream,
+        eof: SequencePointer,
         handle: Handle,
     ) -> VortexResult<LayoutRef> {
         let dtype = stream.dtype().clone();
         let buffer_size = self.buffer_size;
 
-        // We have no choice but to put our final buffers here!
-        // We cannot hold on to sequence ids across iterations of the stream, otherwise we can
-        // cause deadlocks with other columns that are waiting for us to flush.
-        let mut final_flush = eof.split_off();
-
         let buffered_bytes_counter = self.buffered_bytes.clone();
         let buffered_stream = try_stream! {
+            let stream = stream.peekable();
+            pin_mut!(stream);
+
             let mut nbytes = 0u64;
             let mut chunks = VecDeque::new();
 
@@ -68,11 +67,23 @@ impl LayoutStrategy for BufferedStrategy {
                 buffered_bytes_counter.fetch_add(chunk_size, Ordering::Relaxed);
                 chunks.push_back(chunk);
 
+                // 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() {
+                        let chunk_size = chunk.nbytes();
+                        nbytes -= chunk_size;
+                        buffered_bytes_counter.fetch_sub(chunk_size, Ordering::Relaxed);
+                        yield (sequence_ptr.advance(), chunk)
+                    }
+                    break;
+                }
+
                 if nbytes < 2 * buffer_size {
                     continue;
                 };
 
-                // Wait until we're at 2x the buffer size before flushing 1x the buffer size
+                // 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 {
@@ -85,13 +96,6 @@ impl LayoutStrategy for BufferedStrategy {
                     yield (sequence_ptr.advance(), chunk)
                 }
             }
-
-            // Now the input stream has ended, flush everything
-            while let Some(chunk) = chunks.pop_front() {
-                let chunk_size = chunk.nbytes();
-                buffered_bytes_counter.fetch_sub(chunk_size, Ordering::Relaxed);
-                yield (final_flush.advance(), chunk)
-            }
         };
 
         self.child