ordered segments buffer

Author: onursatici

Cargo.lock

@@ -23,6 +23,7 @@ getrandom_v03 = { workspace = true }
 itertools = { workspace = true }
 linked_hash_set = { workspace = true }
 log = { workspace = true }
+parking_lot = { workspace = true }
 moka = { workspace = true, features = ["future"] }
 object_store = { workspace = true, optional = true }
 pin-project-lite = { workspace = true }

vortex-file/Cargo.toml

@@ -1,4 +1,5 @@
 mod cache;
+mod ordered;
 pub(crate) mod writer;
 
 pub use cache::*;

vortex-file/src/segments/mod.rs

@@ -0,0 +1,110 @@
+use std::collections::{BTreeMap, BTreeSet};
+use std::task::Waker;
+
+use vortex_array::aliases::hash_map::HashMap;
+use vortex_buffer::ByteBuffer;
+use vortex_error::{VortexResult, vortex_bail, vortex_err};
+use vortex_layout::segments::SegmentId;
+
+// [start, end)
+#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash, Debug)]
+pub(super) struct Region {
+    pub start: usize,
+    pub end: usize,
+}
+
+impl Default for Region {
+    fn default() -> Self {
+        Self {
+            start: 0,
+            end: usize::MAX,
+        }
+    }
+}
+
+impl Region {
+    pub fn split(self, splits: usize) -> VortexResult<impl Iterator<Item = Self>> {
+        let step = (self.end - self.start) / splits;
+        if step == 0 {
+            vortex_bail!("region space exhausted!");
+        }
+        Ok((self.start..self.end).step_by(step).map(move |start| Self {
+            start,
+            end: start + step,
+        }))
+    }
+}
+
+pub(super) struct OrderedBuffers {
+    data: BTreeMap<usize, Vec<ByteBuffer>>,
+    active_regions: BTreeSet<Region>,
+    wakers: HashMap<Region, Waker>,
+    next_segment_id: SegmentId,
+}
+
+impl Default for OrderedBuffers {
+    fn default() -> Self {
+        Self {
+            data: Default::default(),
+            active_regions: [Region::default()].into(),
+            wakers: Default::default(),
+            next_segment_id: Default::default(),
+        }
+    }
+}
+
+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) {
+                waker.wake_by_ref();
+            }
+        }
+    }
+
+    pub fn split_region(
+        &mut self,
+        region: &Region,
+        splits: usize,
+    ) -> VortexResult<impl Iterator<Item = Region>> {
+        if !self.active_regions.remove(&region) {
+            vortex_bail!("region not active {:?}", region);
+        }
+        Ok(region.split(splits)?.map(|region| {
+            self.active_regions.insert(region);
+            region
+        }))
+    }
+
+    pub fn insert_buffer(&mut self, idx: usize, buffer: Vec<ByteBuffer>) {
+        self.data.insert(idx, buffer);
+    }
+
+    pub fn register_waker(&mut self, region: Region, waker: Waker) {
+        // TODO(os): should this store a Vec<Waker> instead of replacing?
+        self.wakers.insert(region, waker);
+    }
+
+    pub fn first_region(&self) -> VortexResult<Region> {
+        self.active_regions
+            .first()
+            .copied()
+            .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 next_segment_id(&mut self) -> SegmentId {
+        let res = self.next_segment_id;
+        self.next_segment_id = SegmentId::from(*res + 1);
+        res
+    }
+}

vortex-file/src/segments/ordered.rs

@@ -1,8 +1,14 @@
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use parking_lot::Mutex;
 use vortex_buffer::{Alignment, ByteBuffer};
-use vortex_error::{VortexResult, vortex_err};
+use vortex_error::{VortexResult, vortex_bail, vortex_err};
 use vortex_io::VortexWrite;
 use vortex_layout::segments::{SegmentId, SegmentWriter};
 
+use super::ordered::{OrderedBuffers, Region};
 use crate::footer::SegmentSpec;
 
 /// A segment writer that holds buffers in memory until they are flushed by a writer.
@@ -63,3 +69,114 @@ impl BufferedSegmentWriter {
         Ok(())
     }
 }
+
+#[derive(Default)]
+pub struct InOrderSegmentWriter {
+    buffers: Arc<Mutex<OrderedBuffers>>,
+    region: Region,
+    region_offset: usize,
+}
+
+impl InOrderSegmentWriter {
+    pub async fn put(&mut self, data: Vec<ByteBuffer>) -> VortexResult<SegmentId> {
+        let buffer_idx = self.region.start + self.region_offset;
+        if buffer_idx >= self.region.end {
+            vortex_bail!("region space exhausted!");
+        }
+        self.buffers.lock().insert_buffer(buffer_idx, data);
+        self.region_offset += 1;
+        self.next_segment_id_once_active().await
+    }
+
+    pub fn split(self, splits: usize) -> VortexResult<Vec<Self>> {
+        Ok(self
+            .buffers
+            .lock()
+            .split_region(&self.region, splits)?
+            .map(|region| Self {
+                buffers: self.buffers.clone(),
+                region,
+                region_offset: 0,
+            })
+            .collect())
+    }
+
+    async fn next_segment_id_once_active(&self) -> VortexResult<SegmentId> {
+        WaitRegionFuture {
+            buffers: self.buffers.clone(),
+            region: self.region,
+        }
+        .await
+    }
+
+    pub async fn flush<W: VortexWrite>(
+        &mut self,
+        writer: &mut futures::io::Cursor<W>,
+        segment_specs: &mut Vec<SegmentSpec>,
+    ) -> VortexResult<()> {
+        let completed = {
+            let mut guard = self.buffers.lock();
+            guard.completed_buffers()?
+        };
+        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.
+            // Don't worry, in most cases the caller knows what they're doing and will align the
+            // buffers themselves, inserting padding buffers where necessary.
+            let alignment = buffers
+                .first()
+                .map(|buffer| buffer.alignment())
+                .unwrap_or_else(Alignment::none);
+
+            // Add any padding required to align the segment.
+            let offset = writer.position();
+            let padding = offset.next_multiple_of(*alignment as u64) - offset;
+            if padding > 0 {
+                writer
+                    .write_all(ByteBuffer::zeroed(padding as usize))
+                    .await?;
+            }
+            let offset = writer.position();
+
+            for buffer in buffers {
+                writer.write_all(buffer).await?;
+            }
+
+            segment_specs.push(SegmentSpec {
+                offset,
+                length: u32::try_from(writer.position() - offset)
+                    .map_err(|_| vortex_err!("segment length exceeds maximum u32"))?,
+                alignment,
+            });
+        }
+        Ok(())
+    }
+}
+
+impl Drop for InOrderSegmentWriter {
+    fn drop(&mut self) {
+        self.buffers.lock().finish_region(&self.region);
+    }
+}
+
+struct WaitRegionFuture {
+    buffers: Arc<Mutex<OrderedBuffers>>,
+    region: Region,
+}
+
+impl Future for WaitRegionFuture {
+    type Output = VortexResult<SegmentId>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let mut guard = self.buffers.lock();
+        let current_first = match guard.first_region() {
+            Ok(first) => first,
+            Err(e) => return Poll::Ready(Err(e)),
+        };
+        if self.region == current_first {
+            return Poll::Ready(Ok(guard.next_segment_id()));
+        }
+        guard.register_waker(self.region, cx.waker().clone());
+        Poll::Pending
+    }
+}

vortex-file/src/segments/writer.rs

@@ -1,11 +1,14 @@
+use std::sync::Arc;
+
 use futures::StreamExt;
 use vortex_array::ArrayContext;
 use vortex_array::stats::{PRUNING_STATS, Stat};
 use vortex_array::stream::ArrayStream;
 use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
 use vortex_flatbuffers::{FlatBuffer, FlatBufferRoot, WriteFlatBuffer, WriteFlatBufferExt};
-use vortex_io::VortexWrite;
+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};
@@ -21,6 +24,8 @@ pub struct VortexWriteOptions {
     strategy: Box<dyn LayoutStrategy>,
     exclude_dtype: bool,
     file_statistics: Vec<Stat>,
+    io_dispatcher: IoDispatcher,
+    executor: Option<Arc<dyn TaskExecutor>>,
 }
 
 impl Default for VortexWriteOptions {
@@ -29,6 +34,8 @@ impl Default for VortexWriteOptions {
             strategy: Box::new(VortexLayoutStrategy),
             exclude_dtype: false,
             file_statistics: PRUNING_STATS.to_vec(),
+            io_dispatcher: IoDispatcher::shared(),
+            executor: None,
         }
     }
 }