async

Author: onursatici

vortex-file/src/segments/writer.rs

@@ -1,7 +1,9 @@
+#![allow(dead_code)]
 use std::pin::Pin;
 use std::sync::Arc;
 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};
@@ -12,73 +14,16 @@ use super::ordered::{OrderedBuffers, Region};
 use crate::footer::SegmentSpec;
 
 /// A segment writer that holds buffers in memory until they are flushed by a writer.
-#[derive(Default)]
-pub(crate) struct BufferedSegmentWriter {
-    /// A Vec byte buffers for segments
-    segments_buffers: Vec<Vec<ByteBuffer>>,
-    next_id: SegmentId,
-}
-
-impl SegmentWriter for BufferedSegmentWriter {
-    fn put(&mut self, data: &[ByteBuffer]) -> SegmentId {
-        self.segments_buffers.push(data.to_vec());
-        let id = self.next_id;
-        self.next_id = SegmentId::from(*self.next_id + 1);
-        id
-    }
-}
-
-impl BufferedSegmentWriter {
-    /// Flush the segments to the provided async writer.
-    pub async fn flush_async<W: VortexWrite>(
-        &mut self,
-        writer: &mut futures::io::Cursor<W>,
-        segment_specs: &mut Vec<SegmentSpec>,
-    ) -> VortexResult<()> {
-        for buffers in self.segments_buffers.drain(..) {
-            // 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(())
-    }
-}
-
 #[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> {
+#[async_trait]
+impl SegmentWriter for InOrderSegmentWriter {
+    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!");
@@ -87,7 +32,9 @@ impl InOrderSegmentWriter {
         self.region_offset += 1;
         self.next_segment_id_once_active().await
     }
+}
 
+impl InOrderSegmentWriter {
     pub fn split(self, splits: usize) -> VortexResult<Vec<Self>> {
         Ok(self
             .buffers

vortex-file/src/strategy.rs

@@ -3,6 +3,7 @@
 use std::collections::VecDeque;
 use std::sync::Arc;
 
+use async_trait::async_trait;
 use itertools::Itertools;
 use vortex_array::arcref::ArcRef;
 use vortex_array::arrays::ConstantArray;
@@ -130,8 +131,9 @@ struct BtrBlocksCompressedWriter {
     previous_chunk: Option<PreviousCompression>,
 }
 
+#[async_trait]
 impl LayoutWriter for BtrBlocksCompressedWriter {
-    fn push_chunk(
+    async fn push_chunk(
         &mut self,
         segment_writer: &mut dyn SegmentWriter,
         chunk: ArrayRef,
@@ -197,15 +199,17 @@ impl LayoutWriter for BtrBlocksCompressedWriter {
 
         compressed_chunk.statistics().inherit(chunk.statistics());
 
-        self.child.push_chunk(segment_writer, compressed_chunk)
+        self.child
+            .push_chunk(segment_writer, compressed_chunk)
+            .await
     }
 
-    fn flush(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<()> {
-        self.child.flush(segment_writer)
+    async fn flush(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<()> {
+        self.child.flush(segment_writer).await
     }
 
-    fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
-        self.child.finish(segment_writer)
+    async fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
+        self.child.finish(segment_writer).await
     }
 }
 
@@ -234,8 +238,9 @@ struct BufferedWriter {
     child: Box<dyn LayoutWriter>,
 }
 
+#[async_trait]
 impl LayoutWriter for BufferedWriter {
-    fn push_chunk(
+    async fn push_chunk(
         &mut self,
         segment_writer: &mut dyn SegmentWriter,
         chunk: ArrayRef,
@@ -248,7 +253,7 @@ impl LayoutWriter for BufferedWriter {
             while self.nbytes > self.buffer_size {
                 if let Some(chunk) = self.chunks.pop_front() {
                     self.nbytes -= chunk.nbytes() as u64;
-                    self.child.push_chunk(segment_writer, chunk)?;
+                    self.child.push_chunk(segment_writer, chunk).await?;
                 } else {
                     break;
                 }
@@ -257,15 +262,15 @@ impl LayoutWriter for BufferedWriter {
         Ok(())
     }
 
-    fn flush(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<()> {
+    async 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.push_chunk(segment_writer, chunk).await?;
         }
-        self.child.flush(segment_writer)
+        self.child.flush(segment_writer).await
     }
 
-    fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
-        self.child.finish(segment_writer)
+    async fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
+        self.child.finish(segment_writer).await
     }
 }
 

vortex-file/src/writer.rs

@@ -1,3 +1,4 @@
+#![allow(dead_code)]
 use std::sync::Arc;
 
 use futures::StreamExt;
@@ -12,7 +13,7 @@ use vortex_layout::scan::TaskExecutor;
 use vortex_layout::{LayoutContext, LayoutStrategy, LayoutWriter};
 
 use crate::footer::{FileStatistics, FooterFlatBufferWriter, Postscript, PostscriptSegment};
-use crate::segments::writer::BufferedSegmentWriter;
+use crate::segments::writer::InOrderSegmentWriter;
 use crate::strategy::VortexLayoutStrategy;
 use crate::{EOF_SIZE, MAGIC_BYTES, MAX_FOOTER_SIZE, VERSION};
 
@@ -85,30 +86,24 @@ impl VortexWriteOptions {
 
         // Our buffered message writer accumulates messages for each batch so we can flush them
         // into the file.
-        let mut segment_writer = BufferedSegmentWriter::default();
+        let mut segment_writer = InOrderSegmentWriter::default();
         let mut segment_specs = vec![];
 
         // Then write the stream via the root layout
         while let Some(chunk) = stream.next().await {
             let chunk = chunk?;
-            layout_writer.push_chunk(&mut segment_writer, chunk)?;
+            layout_writer.push_chunk(&mut segment_writer, chunk).await?;
             // NOTE(ngates): we could spawn this task and continue to compress the next chunk.
-            segment_writer
-                .flush_async(&mut write, &mut segment_specs)
-                .await?;
+            segment_writer.flush(&mut write, &mut segment_specs).await?;
         }
 
         // Flush the final layout messages into the file
-        layout_writer.flush(&mut segment_writer)?;
-        segment_writer
-            .flush_async(&mut write, &mut segment_specs)
-            .await?;
+        layout_writer.flush(&mut segment_writer).await?;
+        segment_writer.flush(&mut write, &mut segment_specs).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_specs)
-            .await?;
+        let layout = layout_writer.finish(&mut segment_writer).await?;
+        segment_writer.flush(&mut write, &mut segment_specs).await?;
 
         // We write our footer components in order of least likely to be needed to most likely.
         // DType is the least likely to be needed, as many readers may provide this from an

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

@@ -1,5 +1,6 @@
 use std::sync::Arc;
 
+use async_trait::async_trait;
 use vortex_array::arcref::ArcRef;
 use vortex_array::{ArrayContext, ArrayRef};
 use vortex_dtype::DType;
@@ -56,8 +57,9 @@ impl ChunkedLayoutWriter {
     }
 }
 
+#[async_trait]
 impl LayoutWriter for ChunkedLayoutWriter {
-    fn push_chunk(
+    async fn push_chunk(
         &mut self,
         segment_writer: &mut dyn SegmentWriter,
         chunk: ArrayRef,
@@ -70,24 +72,24 @@ impl LayoutWriter for ChunkedLayoutWriter {
             .options
             .chunk_strategy
             .new_writer(&self.ctx, chunk.dtype())?;
-        chunk_writer.push_chunk(segment_writer, chunk)?;
-        chunk_writer.flush(segment_writer)?;
+        chunk_writer.push_chunk(segment_writer, chunk).await?;
+        chunk_writer.flush(segment_writer).await?;
         self.chunks.push(chunk_writer);
 
         Ok(())
     }
 
-    fn flush(&mut self, _segment_writer: &mut dyn SegmentWriter) -> VortexResult<()> {
+    async fn flush(&mut self, _segment_writer: &mut dyn SegmentWriter) -> VortexResult<()> {
         // We flush each chunk as we write it, so there's nothing to do here.
         Ok(())
     }
 
-    fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
+    async fn finish(&mut self, segment_writer: &mut dyn SegmentWriter) -> VortexResult<Layout> {
         // Call finish on each chunk's writer
         let mut children = vec![];
         for writer in self.chunks.iter_mut() {
             // FIXME(ngates): we should try calling finish after each chunk.
-            children.push(writer.finish(segment_writer)?);
+            children.push(writer.finish(segment_writer).await?);
         }
 
         // If there's only one child, there's no point even writing a stats table since

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

@@ -1,3 +1,4 @@
+use async_trait::async_trait;
 use bytes::Bytes;
 use vortex_array::arcref::ArcRef;
 use vortex_array::vtable::EncodingVTable as _;
@@ -71,14 +72,15 @@ struct DelegatingDictLayoutWriter {
     writer: Option<Box<dyn LayoutWriter>>,
 }
 
+#[async_trait]
 impl LayoutWriter for DelegatingDictLayoutWriter {
-    fn push_chunk(
+    async fn push_chunk(
         &mut self,
         segment_writer: &mut dyn crate::segments::SegmentWriter,
         chunk: ArrayRef,
     ) -> VortexResult<()> {
         match self.writer.as_mut() {
-            Some(writer) => writer.push_chunk(segment_writer, chunk),
+            Some(writer) => writer.push_chunk(segment_writer, chunk).await,
             None => {
                 let compressed = BtrBlocksCompressor.compress(&chunk)?;
                 let mut writer = if !compressed.is_encoding(DictEncoding.id()) {
@@ -91,30 +93,30 @@ impl LayoutWriter for DelegatingDictLayoutWriter {
                     )
                     .boxed()
                 };
-                writer.push_chunk(segment_writer, chunk)?;
+                writer.push_chunk(segment_writer, chunk).await?;
                 self.writer = Some(writer);
                 Ok(())
             }
         }
     }
 
-    fn flush(
+    async fn flush(
         &mut self,
         segment_writer: &mut dyn crate::segments::SegmentWriter,
     ) -> VortexResult<()> {
         match self.writer.as_mut() {
             None => vortex_bail!("flush called before push_chunk"),
-            Some(writer) => writer.flush(segment_writer),
+            Some(writer) => writer.flush(segment_writer).await,
         }
     }
 
-    fn finish(
+    async fn finish(
         &mut self,
         segment_writer: &mut dyn crate::segments::SegmentWriter,
     ) -> VortexResult<Layout> {
         match self.writer.as_mut() {
             None => vortex_bail!("finish called before push_chunk"),
-            Some(writer) => writer.finish(segment_writer),
+            Some(writer) => writer.finish(segment_writer).await,
         }
     }
 }