Improve write strategy (#2783)

* Introducing a separate flush stage allows us to ensure stats tables
are written at the end of the file.
* We buffer each column to help keep chunks from the same column next to
each other, within some memory budget. Currently this is naive
per-column memory, not amortized over the size of the schema.

Note: this isn't a break, but we should re-generate S3 files

Author: gatesn

vortex-file/src/strategy.rs

@@ -3,14 +3,14 @@
 use std::sync::Arc;
 
 use vortex_array::arcref::ArcRef;
+use vortex_array::nbytes::NBytes;
 use vortex_array::stats::{PRUNING_STATS, STATS_TO_WRITE};
 use vortex_array::{Array, ArrayContext, ArrayRef};
 use vortex_btrblocks::BtrBlocksCompressor;
 use vortex_dtype::DType;
 use vortex_error::VortexResult;
 use vortex_layout::layouts::chunked::writer::{ChunkedLayoutOptions, ChunkedLayoutWriter};
 use vortex_layout::layouts::flat::FlatLayout;
-use vortex_layout::layouts::flat::writer::FlatLayoutOptions;
 use vortex_layout::layouts::stats::writer::{StatsLayoutOptions, StatsLayoutWriter};
 use vortex_layout::layouts::struct_::writer::StructLayoutWriter;
 use vortex_layout::segments::SegmentWriter;
@@ -32,13 +32,24 @@ impl LayoutStrategy for VortexLayoutStrategy {
             );
         }
 
-        // Otherwise, we finish with compressing the chunks
+        // Leaf arrays are written as flat arrays, above which, we buffer up to 16MB to try to keep
+        // chunks for the same column next to each other, within some reasonable write-time memory
+        // buffering limit.
+        let writer: ArcRef<dyn LayoutStrategy> = ArcRef::new_arc(Arc::new(BufferedStrategy {
+            child: ArcRef::new_ref(&FlatLayout),
+            // Buffer 4MB of compressed data per column before writing the chunks consecutively.
+            // TODO(ngates): this should really be amortized by the number of fields? Maybe the
+            //  strategy could keep track of how many writers were created?
+            buffer_size: 4 << 20, // 4 MB
+        }) as _);
+
+        // Compress each chunk with btrblocks.
         let writer = BtrBlocksCompressedWriter {
             child: ChunkedLayoutWriter::new(
                 ctx.clone(),
                 &DType::Null,
                 ChunkedLayoutOptions {
-                    chunk_strategy: ArcRef::new_arc(Arc::new(FlatLayoutOptions::default()) as _),
+                    chunk_strategy: writer,
                 },
             )
             .boxed(),
@@ -117,6 +128,64 @@ impl LayoutWriter for BtrBlocksCompressedWriter {
         self.child.push_chunk(segment_writer, compressed)
     }
 
+    fn flush(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<()> {
+        self.child.flush(segment_writer)
+    }
+
+    fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
+        self.child.finish(segment_writer)
+    }
+}
+
+struct BufferedStrategy {
+    child: ArcRef<dyn LayoutStrategy>,
+    buffer_size: u64,
+}
+
+impl LayoutStrategy for BufferedStrategy {
+    fn new_writer(&self, ctx: &ArrayContext, dtype: &DType) -> VortexResult<Box<dyn LayoutWriter>> {
+        let child = self.child.new_writer(ctx, dtype)?;
+        Ok(BufferedWriter {
+            chunks: Vec::new(),
+            nbytes: 0,
+            buffer_size: self.buffer_size,
+            child,
+        }
+        .boxed())
+    }
+}
+
+struct BufferedWriter {
+    chunks: Vec<ArrayRef>,
+    nbytes: u64,
+    buffer_size: u64,
+    child: Box<dyn LayoutWriter>,
+}
+
+impl LayoutWriter for BufferedWriter {
+    fn push_chunk(
+        &mut self,
+        segment_writer: &mut dyn SegmentWriter,
+        chunk: ArrayRef,
+    ) -> VortexResult<()> {
+        self.nbytes += chunk.nbytes() as u64;
+        self.chunks.push(chunk);
+        if self.nbytes >= self.buffer_size {
+            for chunk in self.chunks.drain(..) {
+                self.child.push_chunk(segment_writer, chunk)?;
+            }
+            self.nbytes = 0;
+        }
+        Ok(())
+    }
+
+    fn flush(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<()> {
+        for chunk in self.chunks.drain(..) {
+            self.child.push_chunk(segment_writer, chunk)?;
+        }
+        self.child.flush(segment_writer)
+    }
+
     fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
         self.child.finish(segment_writer)
     }

vortex-file/src/writer.rs

@@ -88,6 +88,12 @@ impl VortexWriteOptions {
         }
 
         // Flush the final layout messages into the file
+        layout_writer.flush(&mut segment_writer)?;
+        segment_writer
+            .flush_async(&mut write, &mut segment_map)
+            .await?;
+
+        // Finish the layouts and flush the finishing messages into the file
         let layout = layout_writer.finish(&mut segment_writer)?;
         segment_writer
             .flush_async(&mut write, &mut segment_map)

vortex-layout/src/layouts/chunked/writer.rs

@@ -62,11 +62,16 @@ impl LayoutWriter for ChunkedLayoutWriter {
             .chunk_strategy
             .new_writer(&self.ctx, chunk.dtype())?;
         chunk_writer.push_chunk(segment_writer, chunk)?;
+        chunk_writer.flush(segment_writer)?;
         self.chunks.push(chunk_writer);
 
         Ok(())
     }
 
+    fn flush(&mut self, _segment_writer: &mut dyn SegmentWriter) -> VortexResult<()> {
+        Ok(())
+    }
+
     fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
         // Call finish on each chunk's writer
         let mut children = vec![];

vortex-layout/src/layouts/flat/writer.rs

@@ -94,6 +94,10 @@ impl LayoutWriter for FlatLayoutWriter {
         Ok(())
     }
 
+    fn flush(&mut self, _segment_writer: &mut dyn SegmentWriter) -> VortexResult<()> {
+        Ok(())
+    }
+
     fn finish(&mut self, _segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
         self.layout
             .take()

vortex-layout/src/layouts/stats/writer.rs

@@ -96,6 +96,10 @@ impl LayoutWriter for StatsLayoutWriter {
         self.child_writer.push_chunk(segment_writer, chunk)
     }
 
+    fn flush(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<()> {
+        self.child_writer.flush(segment_writer)
+    }
+
     fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
         let child = self.child_writer.finish(segment_writer)?;
 

vortex-layout/src/layouts/struct_/writer.rs

@@ -93,6 +93,13 @@ impl LayoutWriter for StructLayoutWriter {
         Ok(())
     }
 
+    fn flush(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<()> {
+        for writer in self.column_strategies.iter_mut() {
+            writer.flush(segment_writer)?;
+        }
+        Ok(())
+    }
+
     fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
         let mut column_layouts = vec![];
         for writer in self.column_strategies.iter_mut() {

vortex-layout/src/stats.rs

@@ -76,6 +76,10 @@ impl LayoutWriter for FileStatsLayoutWriter {
         self.inner.push_chunk(segment_writer, chunk)
     }
 
+    fn flush(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<()> {
+        self.inner.flush(segment_writer)
+    }
+
     fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
         self.inner.finish(segment_writer)
     }

vortex-layout/src/writer.rs

@@ -7,12 +7,17 @@ use crate::segments::SegmentWriter;
 /// A strategy for writing chunks of an array into a layout.
 // [layout writer]
 pub trait LayoutWriter: Send {
+    /// Push a chunk into the layout writer.
     fn push_chunk(
         &mut self,
         segment_writer: &mut dyn SegmentWriter,
         chunk: ArrayRef,
     ) -> VortexResult<()>;
 
+    /// Flush any buffered chunks.
+    fn flush(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<()>;
+
+    /// Write any final data (e.g. stats) and return the finished [`Layout`].
     fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout>;
 }
 // [layout writer]
@@ -33,6 +38,7 @@ pub trait LayoutWriterExt: LayoutWriter {
         chunk: ArrayRef,
     ) -> VortexResult<Layout> {
         self.push_chunk(segment_writer, chunk)?;
+        self.flush(segment_writer)?;
         self.finish(segment_writer)
     }
 
@@ -46,6 +52,7 @@ pub trait LayoutWriterExt: LayoutWriter {
         for chunk in iter.into_iter() {
             self.push_chunk(segment_writer, chunk?)?
         }
+        self.flush(segment_writer)?;
         self.finish(segment_writer)
     }
 }

vortex-layout/src/writers/repartition.rs

@@ -46,7 +46,7 @@ impl RepartitionWriter {
         }
     }
 
-    fn flush(&mut self, segments: &mut dyn SegmentWriter) -> VortexResult<()> {
+    fn maybe_flush_chunk(&mut self, segments: &mut dyn SegmentWriter) -> VortexResult<()> {
         if self.nbytes >= self.options.block_size_minimum {
             let nblocks = self.row_count / self.options.block_len_multiple;
 
@@ -112,18 +112,22 @@ impl LayoutWriter for RepartitionWriter {
             self.chunks.push_back(c);
             offset = end;
 
-            self.flush(segment_writer)?;
+            self.maybe_flush_chunk(segment_writer)?;
         }
 
         Ok(())
     }
 
-    fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
+    fn flush(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<()> {
         let chunk =
             ChunkedArray::new_unchecked(self.chunks.drain(..).collect(), self.dtype.clone())
                 .to_canonical()?
                 .into_array();
         self.writer.push_chunk(segment_writer, chunk)?;
+        self.writer.flush(segment_writer)
+    }
+
+    fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
         self.writer.finish(segment_writer)
     }
 }