fix flush

Author: onursatici

bench-vortex/src/compress/vortex.rs

@@ -7,12 +7,13 @@ use tokio::runtime::Handle;
 use vortex::Array;
 use vortex::arrow::IntoArrowArray;
 use vortex::error::VortexResult;
-use vortex::file::{VortexOpenOptions, VortexWriteOptions};
+use vortex::file::{VortexLayoutStrategy, VortexOpenOptions, VortexWriteOptions};
 use vortex::stream::ArrayStreamArrayExt;
 
 #[inline(never)]
 pub async fn vortex_compress_write(array: &dyn Array, buf: &mut Vec<u8>) -> VortexResult<u64> {
     Ok(VortexWriteOptions::default()
+        .with_strategy(VortexLayoutStrategy::default().with_tokio_executor(Handle::current()))
         .write(Cursor::new(buf), array.to_array_stream())
         .await?
         .position())

encodings/dict/src/builders/mod.rs

@@ -23,7 +23,7 @@ pub const UNCONSTRAINED: DictConstraints = DictConstraints {
     max_len: usize::MAX,
 };
 
-pub trait DictEncoder: Send {
+pub trait DictEncoder: Send + Sync {
     fn encode(&mut self, array: &dyn Array) -> VortexResult<ArrayRef>;
 
     fn values(&mut self) -> VortexResult<ArrayRef>;

vortex-file/src/segments/ordered.rs

@@ -28,10 +28,13 @@ impl Region {
         if step == 0 {
             vortex_bail!("region space exhausted!");
         }
-        Ok((self.start..self.end).step_by(step).map(move |start| Self {
-            start,
-            end: start + step,
-        }))
+        Ok((self.start..self.end)
+            .step_by(step)
+            .skip(1)
+            .map(move |start| Self {
+                start,
+                end: start + step,
+            }))
     }
 }
 
@@ -57,7 +60,7 @@ impl OrderedBuffers {
     pub fn finish_region(&mut self, region: &Region) {
         self.active_regions.remove(&region);
         if let Ok(first) = self.first_region() {
-            if let Some(waker) = self.wakers.get(&first) {
+            if let Some(waker) = self.wakers.remove(&first) {
                 waker.wake_by_ref();
             }
         }
@@ -93,13 +96,15 @@ impl OrderedBuffers {
             .ok_or_else(|| vortex_err!("no active regions"))
     }
 
-    pub fn completed_buffers(&mut self) -> VortexResult<BTreeMap<usize, Vec<ByteBuffer>>> {
-        let completed_until = self.first_region()?;
-        let mut completed = std::mem::take(&mut self.data);
-        // truncate completed to include only completed regions
-        let in_progress = completed.split_off(&completed_until.start);
-        self.data = in_progress;
-        Ok(completed)
+    pub fn take_buffers(&mut self) -> VortexResult<BTreeMap<usize, Vec<ByteBuffer>>> {
+        if self.active_regions.len() > 1 {
+            vortex_bail!("there are more than one active writers");
+        }
+        if !self.wakers.is_empty() {
+            vortex_bail!("there is an inflight write");
+        }
+        self.active_regions = [Region::default()].into();
+        Ok(std::mem::take(&mut self.data))
     }
 
     pub fn next_segment_id(&mut self) -> SegmentId {

vortex-file/src/segments/writer.rs

@@ -6,9 +6,9 @@ use std::task::{Context, Poll};
 use async_trait::async_trait;
 use parking_lot::Mutex;
 use vortex_buffer::{Alignment, ByteBuffer};
-use vortex_error::{VortexResult, vortex_bail, vortex_err};
+use vortex_error::{VortexExpect as _, VortexResult, vortex_bail, vortex_err};
 use vortex_io::VortexWrite;
-use vortex_layout::segments::{SegmentId, SegmentWriter};
+use vortex_layout::segments::{ConcurrentSegmentWriter, SegmentId, SegmentWriter};
 
 use super::ordered::{OrderedBuffers, Region};
 use crate::footer::SegmentSpec;
@@ -34,20 +34,38 @@ impl SegmentWriter for InOrderSegmentWriter {
     }
 }
 
-impl InOrderSegmentWriter {
-    pub fn split(self, splits: usize) -> VortexResult<Vec<Self>> {
-        Ok(self
+impl ConcurrentSegmentWriter for InOrderSegmentWriter {
+    fn split_off(&mut self, splits: usize) -> VortexResult<Vec<Box<dyn ConcurrentSegmentWriter>>> {
+        let unwritten_region = Region {
+            start: self.region.start + self.region_offset,
+            end: self.region.end,
+        };
+        let mut regions: Vec<_> = self
             .buffers
             .lock()
-            .split_region(&self.region, splits)?
-            .map(|region| Self {
-                buffers: self.buffers.clone(),
-                region,
-                region_offset: 0,
+            .split_region(&unwritten_region, splits + 1)?
+            .collect();
+        // assign last splits region to self
+        let last = regions
+            .pop()
+            .vortex_expect("there must be at least 1 split");
+        self.region = last;
+        self.region_offset = 0;
+
+        Ok(regions
+            .into_iter()
+            .map(|region| {
+                Box::new(Self {
+                    buffers: self.buffers.clone(),
+                    region,
+                    region_offset: 0,
+                }) as Box<dyn ConcurrentSegmentWriter>
             })
             .collect())
     }
+}
 
+impl InOrderSegmentWriter {
     async fn next_segment_id_once_active(&self) -> VortexResult<SegmentId> {
         WaitRegionFuture {
             buffers: self.buffers.clone(),
@@ -61,10 +79,12 @@ impl InOrderSegmentWriter {
         writer: &mut futures::io::Cursor<W>,
         segment_specs: &mut Vec<SegmentSpec>,
     ) -> VortexResult<()> {
-        let completed = {
-            let mut guard = self.buffers.lock();
-            guard.completed_buffers()?
-        };
+        let completed = self.buffers.lock().take_buffers()?;
+        // we are the only writer if here, reclaim the entire region
+        self.region = Region::default();
+        self.region_offset = 0;
+
+        // TODO(os): spawn everything below
         for buffers in completed.into_values() {
             // The API requires us to write these buffers contiguously. Therefore, we can only
             // respect the alignment of the first one.

vortex-file/src/strategy.rs

@@ -21,22 +21,37 @@ use vortex_layout::layouts::repartition::{
 };
 use vortex_layout::layouts::stats::writer::{StatsLayoutOptions, StatsLayoutWriter};
 use vortex_layout::layouts::struct_::writer::StructLayoutWriter;
-use vortex_layout::segments::SegmentWriter;
+use vortex_layout::scan::TaskExecutor;
+use vortex_layout::segments::ConcurrentSegmentWriter;
 use vortex_layout::{Layout, LayoutStrategy, LayoutWriter, LayoutWriterExt};
 
 const ROW_BLOCK_SIZE: usize = 8192;
 
 /// The default Vortex file layout strategy.
-#[derive(Clone, Debug, Default)]
-pub struct VortexLayoutStrategy;
+#[derive(Clone, Default)]
+pub struct VortexLayoutStrategy {
+    executor: Option<Arc<dyn TaskExecutor>>,
+}
+
+impl VortexLayoutStrategy {
+    #[cfg(feature = "tokio")]
+    pub fn with_tokio_executor(mut self, handle: tokio::runtime::Handle) -> Self {
+        self.executor = Some(Arc::new(handle));
+        self
+    }
+}
 
 impl LayoutStrategy for VortexLayoutStrategy {
     fn new_writer(&self, ctx: &ArrayContext, dtype: &DType) -> VortexResult<Box<dyn LayoutWriter>> {
         // First, we unwrap struct arrays into their components.
         if dtype.is_struct() {
-            return Ok(
-                StructLayoutWriter::try_new_with_strategy(ctx, dtype, self.clone())?.boxed(),
-            );
+            return Ok(StructLayoutWriter::try_new_with_strategy(
+                ctx,
+                dtype,
+                self.executor.clone(),
+                self.clone(),
+            )?
+            .boxed());
         }
 
         // We buffer arrays per column, before flushing them into a chunked layout.
@@ -135,7 +150,7 @@ struct BtrBlocksCompressedWriter {
 impl LayoutWriter for BtrBlocksCompressedWriter {
     async fn push_chunk(
         &mut self,
-        segment_writer: &mut dyn SegmentWriter,
+        segment_writer: &mut dyn ConcurrentSegmentWriter,
         chunk: ArrayRef,
     ) -> VortexResult<()> {
         // Compute the stats for the chunk prior to compression
@@ -204,11 +219,17 @@ impl LayoutWriter for BtrBlocksCompressedWriter {
             .await
     }
 
-    async fn flush(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<()> {
+    async fn flush(
+        &mut self,
+        segment_writer: &mut dyn ConcurrentSegmentWriter,
+    ) -> VortexResult<()> {
         self.child.flush(segment_writer).await
     }
 
-    async fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
+    async fn finish(
+        &mut self,
+        segment_writer: &mut dyn ConcurrentSegmentWriter,
+    ) -> VortexResult<Layout> {
         self.child.finish(segment_writer).await
     }
 }
@@ -242,7 +263,7 @@ struct BufferedWriter {
 impl LayoutWriter for BufferedWriter {
     async fn push_chunk(
         &mut self,
-        segment_writer: &mut dyn SegmentWriter,
+        segment_writer: &mut dyn ConcurrentSegmentWriter,
         chunk: ArrayRef,
     ) -> VortexResult<()> {
         self.nbytes += chunk.nbytes() as u64;
@@ -262,14 +283,20 @@ impl LayoutWriter for BufferedWriter {
         Ok(())
     }
 
-    async fn flush(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<()> {
+    async fn flush(
+        &mut self,
+        segment_writer: &mut dyn ConcurrentSegmentWriter,
+    ) -> VortexResult<()> {
         for chunk in self.chunks.drain(..) {
             self.child.push_chunk(segment_writer, chunk).await?;
         }
         self.child.flush(segment_writer).await
     }
 
-    async fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
+    async fn finish(
+        &mut self,
+        segment_writer: &mut dyn ConcurrentSegmentWriter,
+    ) -> VortexResult<Layout> {
         self.child.finish(segment_writer).await
     }
 }

vortex-file/src/writer.rs

@@ -1,5 +1,4 @@
 #![allow(dead_code)]
-use std::sync::Arc;
 
 use futures::StreamExt;
 use vortex_array::ArrayContext;
@@ -9,7 +8,6 @@ use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
 use vortex_flatbuffers::{FlatBuffer, FlatBufferRoot, WriteFlatBuffer, WriteFlatBufferExt};
 use vortex_io::{IoDispatcher, VortexWrite};
 use vortex_layout::layouts::file_stats::FileStatsLayoutWriter;
-use vortex_layout::scan::TaskExecutor;
 use vortex_layout::{LayoutContext, LayoutStrategy, LayoutWriter};
 
 use crate::footer::{FileStatistics, FooterFlatBufferWriter, Postscript, PostscriptSegment};
@@ -26,17 +24,15 @@ pub struct VortexWriteOptions {
     exclude_dtype: bool,
     file_statistics: Vec<Stat>,
     io_dispatcher: IoDispatcher,
-    executor: Option<Arc<dyn TaskExecutor>>,
 }
 
 impl Default for VortexWriteOptions {
     fn default() -> Self {
         Self {
-            strategy: Box::new(VortexLayoutStrategy),
+            strategy: Box::new(VortexLayoutStrategy::default()),
             exclude_dtype: false,
             file_statistics: PRUNING_STATS.to_vec(),
             io_dispatcher: IoDispatcher::shared(),
-            executor: None,
         }
     }
 }

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

@@ -9,7 +9,7 @@ use vortex_error::{VortexExpect, VortexResult};
 use crate::data::Layout;
 use crate::layouts::chunked::ChunkedLayout;
 use crate::layouts::flat::writer::FlatLayoutStrategy;
-use crate::segments::SegmentWriter;
+use crate::segments::ConcurrentSegmentWriter;
 use crate::strategy::LayoutStrategy;
 use crate::writer::LayoutWriter;
 use crate::{LayoutVTableRef, LayoutWriterExt};
@@ -61,7 +61,7 @@ impl ChunkedLayoutWriter {
 impl LayoutWriter for ChunkedLayoutWriter {
     async fn push_chunk(
         &mut self,
-        segment_writer: &mut dyn SegmentWriter,
+        segment_writer: &mut dyn ConcurrentSegmentWriter,
         chunk: ArrayRef,
     ) -> VortexResult<()> {
         self.row_count += chunk.len() as u64;
@@ -79,12 +79,18 @@ impl LayoutWriter for ChunkedLayoutWriter {
         Ok(())
     }
 
-    async fn flush(&mut self, _segment_writer: &mut dyn SegmentWriter) -> VortexResult<()> {
+    async fn flush(
+        &mut self,
+        _segment_writer: &mut dyn ConcurrentSegmentWriter,
+    ) -> VortexResult<()> {
         // We flush each chunk as we write it, so there's nothing to do here.
         Ok(())
     }
 
-    async fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
+    async fn finish(
+        &mut self,
+        segment_writer: &mut dyn ConcurrentSegmentWriter,
+    ) -> VortexResult<Layout> {
         // Call finish on each chunk's writer
         let mut children = vec![];
         for writer in self.chunks.iter_mut() {

vortex-layout/src/layouts/dict/writer/mod.rs

@@ -76,7 +76,7 @@ struct DelegatingDictLayoutWriter {
 impl LayoutWriter for DelegatingDictLayoutWriter {
     async fn push_chunk(
         &mut self,
-        segment_writer: &mut dyn crate::segments::SegmentWriter,
+        segment_writer: &mut dyn crate::segments::ConcurrentSegmentWriter,
         chunk: ArrayRef,
     ) -> VortexResult<()> {
         match self.writer.as_mut() {
@@ -102,7 +102,7 @@ impl LayoutWriter for DelegatingDictLayoutWriter {
 
     async fn flush(
         &mut self,
-        segment_writer: &mut dyn crate::segments::SegmentWriter,
+        segment_writer: &mut dyn crate::segments::ConcurrentSegmentWriter,
     ) -> VortexResult<()> {
         match self.writer.as_mut() {
             None => vortex_bail!("flush called before push_chunk"),
@@ -112,7 +112,7 @@ impl LayoutWriter for DelegatingDictLayoutWriter {
 
     async fn finish(
         &mut self,
-        segment_writer: &mut dyn crate::segments::SegmentWriter,
+        segment_writer: &mut dyn crate::segments::ConcurrentSegmentWriter,
     ) -> VortexResult<Layout> {
         match self.writer.as_mut() {
             None => vortex_bail!("finish called before push_chunk"),