layout writer tweaks (#3435)

Signed-off-by: Onur Satici <[email protected]>

Author: onursatici

vortex-array/src/context.rs

@@ -50,7 +50,7 @@ impl ArrayRegistry {
 #[derive(Debug, Clone)]
 pub struct VTableContext<T>(Arc<RwLock<Vec<T>>>);
 
-impl<T: Clone + Eq + Display> VTableContext<T> {
+impl<T: Clone + Eq> VTableContext<T> {
     pub fn empty() -> Self {
         Self(Arc::new(RwLock::new(Vec::new())))
     }

vortex-file/src/strategy.rs

@@ -1,26 +1,19 @@
 //! This module defines the default layout strategy for a Vortex file.
 
-use std::collections::VecDeque;
 use std::sync::Arc;
 
 use arcref::ArcRef;
-use async_stream::try_stream;
-use futures::{FutureExt, StreamExt, pin_mut};
-use vortex_array::stats::{PRUNING_STATS, Stat};
-use vortex_array::{Array, ArrayContext};
-use vortex_btrblocks::BtrBlocksCompressor;
+use vortex_array::stats::PRUNING_STATS;
+use vortex_layout::LayoutStrategy;
+use vortex_layout::layouts::buffered::BufferedStrategy;
 use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy;
+use vortex_layout::layouts::compressed::BtrBlocksCompressedStrategy;
 use vortex_layout::layouts::dict::writer::DictStrategy;
 use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
 use vortex_layout::layouts::repartition::{RepartitionStrategy, RepartitionWriterOptions};
 use vortex_layout::layouts::struct_::writer::StructStrategy;
 use vortex_layout::layouts::zoned::writer::{ZonedLayoutOptions, ZonedStrategy};
-use vortex_layout::scan::{TaskExecutor, TaskExecutorExt};
-use vortex_layout::segments::SequenceWriter;
-use vortex_layout::{
-    LayoutStrategy, SendableLayoutWriter, SendableSequentialStream, SequentialStreamAdapter,
-    SequentialStreamExt as _,
-};
+use vortex_layout::scan::TaskExecutor;
 
 const ROW_BLOCK_SIZE: usize = 8192;
 
@@ -93,121 +86,3 @@ impl VortexLayoutStrategy {
 fn arcref(item: impl LayoutStrategy) -> ArcRef<dyn LayoutStrategy> {
     ArcRef::new_arc(Arc::new(item))
 }
-
-/// A layout writer that compresses chunks using a sampling compressor.
-struct BtrBlocksCompressedStrategy {
-    child: ArcRef<dyn LayoutStrategy>,
-    executor: Arc<dyn TaskExecutor>,
-    parallelism: usize,
-}
-
-impl BtrBlocksCompressedStrategy {
-    pub fn new(
-        child: ArcRef<dyn LayoutStrategy>,
-        executor: Arc<dyn TaskExecutor>,
-        parallelism: usize,
-    ) -> Self {
-        Self {
-            child,
-            executor,
-            parallelism,
-        }
-    }
-}
-
-impl LayoutStrategy for BtrBlocksCompressedStrategy {
-    fn write_stream(
-        &self,
-        ctx: &ArrayContext,
-        sequence_writer: SequenceWriter,
-        stream: SendableSequentialStream,
-    ) -> SendableLayoutWriter {
-        let executor = self.executor.clone();
-
-        let dtype = stream.dtype().clone();
-        let stream = stream
-            .map(|chunk| {
-                async {
-                    let (sequence_id, chunk) = chunk?;
-                    // Compute the stats for the chunk prior to compression
-                    chunk
-                        .statistics()
-                        .compute_all(&Stat::all().collect::<Vec<_>>())?;
-                    Ok((sequence_id, BtrBlocksCompressor.compress(&chunk)?))
-                }
-                .boxed()
-            })
-            .map(move |compress_future| executor.spawn(compress_future))
-            .buffered(self.parallelism);
-
-        self.child.write_stream(
-            ctx,
-            sequence_writer,
-            SequentialStreamAdapter::new(dtype, stream).sendable(),
-        )
-    }
-}
-
-struct BufferedStrategy {
-    child: ArcRef<dyn LayoutStrategy>,
-    buffer_size: u64,
-}
-
-impl BufferedStrategy {
-    pub fn new(child: ArcRef<dyn LayoutStrategy>, buffer_size: u64) -> Self {
-        Self { child, buffer_size }
-    }
-}
-
-impl LayoutStrategy for BufferedStrategy {
-    fn write_stream(
-        &self,
-        ctx: &ArrayContext,
-        sequence_writer: SequenceWriter,
-        stream: SendableSequentialStream,
-    ) -> SendableLayoutWriter {
-        let dtype = stream.dtype().clone();
-        let buffer_size = self.buffer_size;
-        let buffered_stream = try_stream! {
-            let stream = stream.peekable();
-            pin_mut!(stream);
-
-            let mut nbytes = 0u64;
-            let mut chunks = VecDeque::new();
-
-            while let Some(chunk) = stream.as_mut().next().await {
-                let (sequence_id, chunk) = chunk?;
-                nbytes += chunk.nbytes() as u64;
-                chunks.push_back(chunk);
-
-                // if this is the last element, flush everything
-                if stream.as_mut().peek().await.is_none() {
-                    let mut sequence_pointer = sequence_id.descend();
-                    while let Some(chunk) = chunks.pop_front() {
-                        yield (sequence_pointer.advance(), chunk)
-                    }
-                    break;
-                }
-
-                if nbytes < 2 * buffer_size {
-                    continue;
-                };
-                // 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_pointer = sequence_id.descend();
-                while nbytes > buffer_size {
-                    let Some(chunk) = chunks.pop_front() else {
-                        break;
-                    };
-                    nbytes -= chunk.nbytes() as u64;
-                    yield (sequence_pointer.advance(), chunk)
-                }
-            }
-        };
-        self.child.write_stream(
-            ctx,
-            sequence_writer,
-            SequentialStreamAdapter::new(dtype, buffered_stream).sendable(),
-        )
-    }
-}

vortex-file/src/writer.rs

@@ -2,7 +2,7 @@ use std::future;
 use std::sync::Arc;
 
 use arcref::ArcRef;
-use futures::StreamExt;
+use futures::TryStreamExt;
 use futures::future::try_join;
 use vortex_array::ArrayContext;
 use vortex_array::stats::{PRUNING_STATS, Stat};
@@ -33,7 +33,7 @@ pub struct VortexWriteOptions {
 impl Default for VortexWriteOptions {
     fn default() -> Self {
         Self {
-            strategy: VortexLayoutStrategy::with_executor(Arc::new(LocalExecutor {})),
+            strategy: VortexLayoutStrategy::with_executor(Arc::new(LocalExecutor)),
             exclude_dtype: false,
             file_statistics: PRUNING_STATS.to_vec(),
             max_variable_length_statistics_size: 64,
@@ -77,10 +77,8 @@ impl VortexWriteOptions {
         let (segment_writer, flusher) = SerialSegmentWriter::create();
         let sequence_writer = SequenceWriter::new(Box::new(segment_writer));
 
-        let stream = stream.filter(|item| match item {
-            Ok(chunk) => future::ready(!chunk.is_empty()),
-            Err(_) => future::ready(true),
-        });
+        let stream = stream.try_filter(|chunk| future::ready(!chunk.is_empty()));
+
         let stream = sequence_writer.new_sequential(ArrayStreamExt::boxed(
             ArrayStreamAdapter::new(dtype.clone(), stream),
         ));
@@ -95,7 +93,7 @@ impl VortexWriteOptions {
         let mut write = futures::io::Cursor::new(write);
         write.write_all(MAGIC_BYTES).await?;
 
-        let io_fut = async { flusher.flush(write).await };
+        let io_fut = flusher.flush(write);
         let compute_fut = self.strategy.write_stream(&ctx, sequence_writer, stream);
         let (layout, (mut write, segment_specs)) = try_join(compute_fut, io_fut).await?;
 

vortex-layout/src/layouts/buffered.rs

@@ -0,0 +1,76 @@
+use std::collections::VecDeque;
+
+use arcref::ArcRef;
+use async_stream::try_stream;
+use futures::{StreamExt as _, pin_mut};
+use vortex_array::ArrayContext;
+
+use crate::segments::SequenceWriter;
+use crate::{
+    LayoutStrategy, SendableLayoutWriter, SendableSequentialStream, SequentialStreamAdapter,
+    SequentialStreamExt as _,
+};
+
+pub struct BufferedStrategy {
+    child: ArcRef<dyn LayoutStrategy>,
+    buffer_size: u64,
+}
+
+impl BufferedStrategy {
+    pub fn new(child: ArcRef<dyn LayoutStrategy>, buffer_size: u64) -> Self {
+        Self { child, buffer_size }
+    }
+}
+
+impl LayoutStrategy for BufferedStrategy {
+    fn write_stream(
+        &self,
+        ctx: &ArrayContext,
+        sequence_writer: SequenceWriter,
+        stream: SendableSequentialStream,
+    ) -> SendableLayoutWriter {
+        let dtype = stream.dtype().clone();
+        let buffer_size = self.buffer_size;
+        let buffered_stream = try_stream! {
+            let stream = stream.peekable();
+            pin_mut!(stream);
+
+            let mut nbytes = 0u64;
+            let mut chunks = VecDeque::new();
+
+            while let Some(chunk) = stream.as_mut().next().await {
+                let (sequence_id, chunk) = chunk?;
+                nbytes += chunk.nbytes() as u64;
+                chunks.push_back(chunk);
+
+                // if this is the last element, flush everything
+                if stream.as_mut().peek().await.is_none() {
+                    let mut sequence_pointer = sequence_id.descend();
+                    while let Some(chunk) = chunks.pop_front() {
+                        yield (sequence_pointer.advance(), chunk)
+                    }
+                    break;
+                }
+
+                if nbytes < 2 * buffer_size {
+                    continue;
+                };
+                // 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_pointer = sequence_id.descend();
+                while nbytes > buffer_size {
+                    let Some(chunk) = chunks.pop_front() else {
+                        break;
+                    };
+                    nbytes -= chunk.nbytes() as u64;
+                    yield (sequence_pointer.advance(), chunk)
+                }
+            }
+        };
+        self.child.write_stream(
+            ctx,
+            sequence_writer,
+            SequentialStreamAdapter::new(dtype, buffered_stream).sendable(),
+        )
+    }
+}

vortex-layout/src/layouts/compressed.rs

@@ -0,0 +1,68 @@
+use std::sync::Arc;
+
+use arcref::ArcRef;
+use futures::{FutureExt as _, StreamExt as _};
+use vortex_array::ArrayContext;
+use vortex_array::stats::Stat;
+use vortex_btrblocks::BtrBlocksCompressor;
+
+use crate::scan::{TaskExecutor, TaskExecutorExt as _};
+use crate::segments::SequenceWriter;
+use crate::{
+    LayoutStrategy, SendableLayoutWriter, SendableSequentialStream, SequentialStreamAdapter,
+    SequentialStreamExt as _,
+};
+
+/// A layout writer that compresses chunks using a sampling compressor.
+pub struct BtrBlocksCompressedStrategy {
+    child: ArcRef<dyn LayoutStrategy>,
+    executor: Arc<dyn TaskExecutor>,
+    parallelism: usize,
+}
+
+impl BtrBlocksCompressedStrategy {
+    pub fn new(
+        child: ArcRef<dyn LayoutStrategy>,
+        executor: Arc<dyn TaskExecutor>,
+        parallelism: usize,
+    ) -> Self {
+        Self {
+            child,
+            executor,
+            parallelism,
+        }
+    }
+}
+
+impl LayoutStrategy for BtrBlocksCompressedStrategy {
+    fn write_stream(
+        &self,
+        ctx: &ArrayContext,
+        sequence_writer: SequenceWriter,
+        stream: SendableSequentialStream,
+    ) -> SendableLayoutWriter {
+        let executor = self.executor.clone();
+
+        let dtype = stream.dtype().clone();
+        let stream = stream
+            .map(|chunk| {
+                async {
+                    let (sequence_id, chunk) = chunk?;
+                    // Compute the stats for the chunk prior to compression
+                    chunk
+                        .statistics()
+                        .compute_all(&Stat::all().collect::<Vec<_>>())?;
+                    Ok((sequence_id, BtrBlocksCompressor.compress(&chunk)?))
+                }
+                .boxed()
+            })
+            .map(move |compress_future| executor.spawn(compress_future))
+            .buffered(self.parallelism);
+
+        self.child.write_stream(
+            ctx,
+            sequence_writer,
+            SequentialStreamAdapter::new(dtype, stream).sendable(),
+        )
+    }
+}

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

@@ -11,7 +11,7 @@ use vortex_btrblocks::BtrBlocksCompressor;
 use vortex_dict::DictEncoding;
 use vortex_dict::builders::{DictConstraints, DictEncoder, dict_encoder};
 use vortex_dtype::{DType, PType};
-use vortex_error::{VortexExpect, VortexResult, VortexUnwrap, vortex_err};
+use vortex_error::{VortexExpect, VortexResult, VortexUnwrap, vortex_bail, vortex_err};
 
 use super::DictLayout;
 use crate::layouts::chunked::ChunkedLayout;
@@ -203,10 +203,8 @@ impl DictStreamState {
         labeler: &mut DictChunkLabeler,
         chunk: ArrayRef,
     ) -> Vec<VortexResult<DictionaryChunk>> {
-        match self.try_encode(labeler, chunk) {
-            Ok(chunks) => chunks,
-            Err(e) => vec![Err(e)],
-        }
+        self.try_encode(labeler, chunk)
+            .unwrap_or_else(|e| vec![Err(e)])
     }
 
     fn try_encode(
@@ -304,10 +302,9 @@ impl DictEncodedRuns {
 
         let (values_tx, values_rx) = oneshot::channel();
         let values_future = async {
-            match values_rx.await {
-                Ok(values) => values,
-                Err(_) => Err(vortex_err!("sender dropped")),
-            }
+            values_rx
+                .await
+                .unwrap_or_else(|_| vortex_bail!("sender dropped"))
         };
 
         let codes_stream = DictEncodedRunStream {

vortex-layout/src/layouts/mod.rs

@@ -4,7 +4,9 @@ use futures::future::{BoxFuture, Shared};
 use vortex_array::ArrayRef;
 use vortex_error::SharedVortexResult;
 
+pub mod buffered;
 pub mod chunked;
+pub mod compressed;
 pub mod dict;
 pub mod file_stats;
 pub mod filter;