feat: Spawn layout evaluation (#2348)

Spawn split evaluation on a dedicated runtime --------- Co-authored-by: Nicholas Gates <[email protected]>
spiraldb · Feb 18, 2025 · 661f6cd · 661f6cd
1 parent 80ec2f6
commit 661f6cd
Show file tree

Hide file tree

Showing 26 changed files with 338 additions and 364 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/vortex-datafusion/Cargo.toml b/vortex-datafusion/Cargo.toml
@@ -47,7 +47,7 @@ vortex-error = { workspace = true, features = ["datafusion"] }
 vortex-expr = { workspace = true, features = ["datafusion"] }
 vortex-file = { workspace = true, features = ["object_store", "tokio"] }
 vortex-io = { workspace = true, features = ["object_store", "tokio"] }
-vortex-layout = { workspace = true }
+vortex-layout = { workspace = true, features = ["tokio"] }
 [features]
 tracing = ["dep:tracing", "dep:tracing-futures", "vortex-io/tracing"]
 

diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs
@@ -5,9 +5,11 @@ use datafusion::datasource::physical_plan::{FileMeta, FileOpenFuture, FileOpener
 use datafusion_common::Result as DFResult;
 use futures::{FutureExt as _, StreamExt};
 use object_store::{ObjectStore, ObjectStoreScheme};
+use tokio::runtime::Handle;
 use vortex_array::{ContextRef, IntoArrayVariant};
 use vortex_error::VortexResult;
 use vortex_expr::{ExprRef, VortexExpr};
+use vortex_file::executor::{TaskExecutor, TokioExecutor};
 use vortex_file::{SplitBy, VortexOpenOptions};
 use vortex_io::ObjectStoreReadAt;
 
@@ -65,6 +67,7 @@ impl FileOpener for VortexFileOpener {
         let object_store = self.object_store.clone();
         let projected_arrow_schema = self.projected_arrow_schema.clone();
         let batch_size = self.batch_size;
+        let executor = TaskExecutor::Tokio(TokioExecutor::new(Handle::current()));
 
         Ok(async move {
             let vxf = VortexOpenOptions::file(read_at)
@@ -86,6 +89,7 @@ impl FileOpener for VortexFileOpener {
                 // but at the moment our scanner has too much overhead to process small
                 // batches efficiently.
                 .with_split_by(SplitBy::RowCount(8 * batch_size))
+                .with_task_executor(executor)
                 .into_array_stream()?
                 .map(move |array| {
                     let st = array?.into_struct()?;

diff --git a/vortex-file/src/exec/inline.rs b/vortex-file/src/exec/inline.rs
diff --git a/vortex-file/src/exec/mod.rs b/vortex-file/src/exec/mod.rs
diff --git a/vortex-file/src/exec/mode.rs b/vortex-file/src/exec/mode.rs
diff --git a/vortex-file/src/exec/tokio.rs b/vortex-file/src/exec/tokio.rs
diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs
@@ -86,7 +86,6 @@
 //! buffers, and [cloud storage](vortex_io::ObjectStoreReadAt), can be used as the "linear and
 //! contiguous memory".
 
-pub mod exec;
 mod file;
 mod footer;
 mod generic;

diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml
@@ -21,18 +21,20 @@ bit-vec = { workspace = true }
 bytes = { workspace = true }
 exponential-decay-histogram = { workspace = true }
 flatbuffers = { workspace = true }
-futures = { workspace = true, features = ["alloc"] }
+flume = { workspace = true }
+futures = { workspace = true, features = ["alloc", "executor"] }
 itertools = { workspace = true }
 log = { workspace = true }
 pin-project-lite = { workspace = true }
 vortex-array = { workspace = true }
 vortex-buffer = { workspace = true }
 vortex-dtype = { workspace = true }
-vortex-error = { workspace = true }
+vortex-error = { workspace = true, features = ["tokio"] }
 vortex-expr = { workspace = true }
 vortex-flatbuffers = { workspace = true, features = ["layout"] }
 vortex-mask = { workspace = true }
 vortex-scalar = { workspace = true }
+tokio = { workspace = true, optional = true }
 
 [dev-dependencies]
 futures = { workspace = true, features = ["executor"] }

diff --git a/vortex-layout/src/data.rs b/vortex-layout/src/data.rs
@@ -12,8 +12,7 @@ use vortex_flatbuffers::{layout as fb, layout, FlatBufferRoot, WriteFlatBuffer};
 
 use crate::context::LayoutContextRef;
 use crate::reader::LayoutReader;
-use crate::scan::ScanExecutor;
-use crate::segments::SegmentId;
+use crate::segments::{AsyncSegmentReader, SegmentId};
 use crate::vtable::LayoutVTableRef;
 use crate::LayoutId;
 
@@ -292,10 +291,10 @@ impl Layout {
     /// Create a reader for this layout.
     pub fn reader(
         &self,
-        executor: Arc<ScanExecutor>,
+        segment_reader: Arc<dyn AsyncSegmentReader>,
         ctx: ContextRef,
     ) -> VortexResult<Arc<dyn LayoutReader + 'static>> {
-        self.encoding().reader(self.clone(), ctx, executor)
+        self.encoding().reader(self.clone(), ctx, segment_reader)
     }
 
     /// Register splits for this layout.

diff --git a/vortex-layout/src/layouts/chunked/eval_expr.rs b/vortex-layout/src/layouts/chunked/eval_expr.rs
@@ -145,6 +145,7 @@ mod test {
     use std::sync::Arc;
 
     use futures::executor::block_on;
+    use rstest::{fixture, rstest};
     use vortex_array::array::{BoolArray, ChunkedArray, ConstantArray};
     use vortex_array::{IntoArray, IntoArrayVariant};
     use vortex_buffer::buffer;
@@ -154,13 +155,14 @@ mod test {
     use vortex_expr::{gt, lit, Identity};
 
     use crate::layouts::chunked::writer::ChunkedLayoutWriter;
-    use crate::scan::ScanExecutor;
     use crate::segments::test::TestSegments;
+    use crate::segments::AsyncSegmentReader;
     use crate::writer::LayoutWriterExt;
     use crate::{Layout, RowMask};
 
+    #[fixture]
     /// Create a chunked layout with three chunks of primitive arrays.
-    fn chunked_layout() -> (Arc<ScanExecutor>, Layout) {
+    fn chunked_layout() -> (Arc<dyn AsyncSegmentReader>, Layout) {
         let mut segments = TestSegments::default();
         let layout = ChunkedLayoutWriter::new(
             &DType::Primitive(PType::I32, NonNullable),
@@ -175,14 +177,14 @@ mod test {
             ],
         )
         .unwrap();
-        (ScanExecutor::inline(Arc::new(segments)), layout)
+        (Arc::new(segments), layout)
     }
 
-    #[test]
-    fn test_chunked_evaluator() {
+    #[rstest]
+    fn test_chunked_evaluator(
+        #[from(chunked_layout)] (segments, layout): (Arc<dyn AsyncSegmentReader>, Layout),
+    ) {
         block_on(async {
-            let (segments, layout) = chunked_layout();
-
             let result = layout
                 .reader(segments, Default::default())
                 .unwrap()
@@ -200,10 +202,11 @@ mod test {
         })
     }
 
-    #[test]
-    fn test_chunked_pruning_mask() {
+    #[rstest]
+    fn test_chunked_pruning_mask(
+        #[from(chunked_layout)] (segments, layout): (Arc<dyn AsyncSegmentReader>, Layout),
+    ) {
         block_on(async {
-            let (segments, layout) = chunked_layout();
             let row_count = layout.row_count();
             let reader = layout.reader(segments, Default::default()).unwrap();
 

diff --git a/vortex-layout/src/layouts/chunked/mod.rs b/vortex-layout/src/layouts/chunked/mod.rs
@@ -13,7 +13,7 @@ use vortex_error::VortexResult;
 use crate::data::Layout;
 use crate::layouts::chunked::reader::ChunkedReader;
 use crate::reader::{LayoutReader, LayoutReaderExt};
-use crate::scan::ScanExecutor;
+use crate::segments::AsyncSegmentReader;
 use crate::vtable::LayoutVTable;
 use crate::{LayoutId, CHUNKED_LAYOUT_ID};
 
@@ -33,9 +33,9 @@ impl LayoutVTable for ChunkedLayout {
         &self,
         layout: Layout,
         ctx: ContextRef,
-        executor: Arc<ScanExecutor>,
+        segment_reader: Arc<dyn AsyncSegmentReader>,
     ) -> VortexResult<Arc<dyn LayoutReader>> {
-        Ok(ChunkedReader::try_new(layout, ctx, executor)?.into_arc())
+        Ok(ChunkedReader::try_new(layout, ctx, segment_reader)?.into_arc())
     }
 
     fn register_splits(

diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs
@@ -14,7 +14,7 @@ use vortex_mask::Mask;
 use crate::layouts::chunked::stats_table::StatsTable;
 use crate::layouts::chunked::ChunkedLayout;
 use crate::reader::LayoutReader;
-use crate::scan::ScanExecutor;
+use crate::segments::AsyncSegmentReader;
 use crate::{ExprEvaluator, Layout, LayoutVTable, RowMask};
 
 type PruningCache = Arc<OnceCell<Option<Mask>>>;
@@ -23,7 +23,7 @@ type PruningCache = Arc<OnceCell<Option<Mask>>>;
 pub struct ChunkedReader {
     layout: Layout,
     ctx: ContextRef,
-    executor: Arc<ScanExecutor>,
+    segment_reader: Arc<dyn AsyncSegmentReader>,
 
     /// A cache of expr -> optional pruning result (applying the pruning expr to the stats table)
     pruning_result: Arc<RwLock<HashMap<ExprRef, PruningCache>>>,
@@ -39,7 +39,7 @@ impl ChunkedReader {
     pub(super) fn try_new(
         layout: Layout,
         ctx: ContextRef,
-        executor: Arc<ScanExecutor>,
+        segment_reader: Arc<dyn AsyncSegmentReader>,
     ) -> VortexResult<Self> {
         if layout.encoding().id() != ChunkedLayout.id() {
             vortex_panic!("Mismatched layout ID")
@@ -71,7 +71,7 @@ impl ChunkedReader {
         Ok(Self {
             layout,
             ctx,
-            executor,
+            segment_reader,
             pruning_result: Arc::new(RwLock::new(HashMap::new())),
             stats_table: Arc::new(OnceCell::new()),
             chunk_readers,
@@ -101,7 +101,7 @@ impl ChunkedReader {
                         let stats_layout = self.layout.child(nchunks, stats_dtype.clone(), "stats")?;
 
                         let stats_array = stats_layout
-                            .reader(self.executor.clone(), self.ctx.clone())?
+                            .reader(self.segment_reader.clone(), self.ctx.clone())?
                             .evaluate_expr(
                                 RowMask::new_valid_between(0, nchunks as u64),
                                 Identity::new_expr(),
@@ -161,7 +161,7 @@ impl ChunkedReader {
             let child_layout =
                 self.layout
                     .child(idx, self.layout.dtype().clone(), format!("[{}]", idx))?;
-            child_layout.reader(self.executor.clone(), self.ctx.clone())
+            child_layout.reader(self.segment_reader.clone(), self.ctx.clone())
         })
     }