feat: Spawn layout evaluation

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"]
 

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()?;

vortex-error/Cargo.toml

@@ -35,6 +35,7 @@ rancor = { workspace = true, optional = true }
 thiserror = { workspace = true }
 url = { workspace = true }
 worker = { workspace = true, optional = true }
+tokio = { workspace = true, optional = true, features = ["rt"] }
 
 [lints]
 workspace = true

vortex-error/src/lib.rs

@@ -216,6 +216,14 @@ pub enum VortexError {
         #[backtrace]
         TryFromIntError,
     ),
+    /// A wrapper for Tokio join error.
+    #[cfg(feature = "tokio")]
+    #[error(transparent)]
+    JoinError(
+        #[from]
+        #[backtrace]
+        tokio::task::JoinError,
+    ),
 }
 
 impl VortexError {

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"] }

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

@@ -175,7 +175,7 @@ mod test {
             ],
         )
         .unwrap();
-        (ScanExecutor::inline(Arc::new(segments)), layout)
+        (ScanExecutor::new(Arc::new(segments)), layout)
     }
 
     #[test]

vortex-layout/src/layouts/flat/eval_expr.rs

@@ -35,7 +35,7 @@ impl ExprEvaluator for FlatReader {
             tasks.push(ScanTask::Expr(expr));
         }
 
-        self.executor().evaluate(&array, &tasks).await
+        self.executor().evaluate(&array, &tasks)
     }
 
     async fn prune_mask(&self, row_mask: RowMask, _expr: ExprRef) -> VortexResult<RowMask> {
@@ -72,7 +72,7 @@ mod test {
                 .unwrap();
 
             let result = layout
-                .reader(ScanExecutor::inline(Arc::new(segments)), Default::default())
+                .reader(ScanExecutor::new(Arc::new(segments)), Default::default())
                 .unwrap()
                 .evaluate_expr(
                     RowMask::new_valid_between(0, layout.row_count()),
@@ -98,7 +98,7 @@ mod test {
 
             let expr = gt(Identity::new_expr(), lit(3i32));
             let result = layout
-                .reader(ScanExecutor::inline(Arc::new(segments)), Default::default())
+                .reader(ScanExecutor::new(Arc::new(segments)), Default::default())
                 .unwrap()
                 .evaluate_expr(RowMask::new_valid_between(0, layout.row_count()), expr)
                 .await
@@ -123,7 +123,7 @@ mod test {
                 .unwrap();
 
             let result = layout
-                .reader(ScanExecutor::inline(Arc::new(segments)), Default::default())
+                .reader(ScanExecutor::new(Arc::new(segments)), Default::default())
                 .unwrap()
                 .evaluate_expr(RowMask::new_valid_between(2, 4), ident())
                 .await

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

@@ -119,7 +119,7 @@ mod tests {
                 .unwrap();
 
             let result = layout
-                .reader(ScanExecutor::inline(Arc::new(segments)), Default::default())
+                .reader(ScanExecutor::new(Arc::new(segments)), Default::default())
                 .unwrap()
                 .evaluate_expr(RowMask::new_valid_between(0, layout.row_count()), ident())
                 .await

vortex-layout/src/layouts/struct_/eval_expr.rs

@@ -54,7 +54,6 @@ impl ExprEvaluator for StructReader {
 
         self.executor()
             .evaluate(&root_scope, &[ScanTask::Expr(partitioned.root.clone())])
-            .await
     }
 
     async fn prune_mask(&self, row_mask: RowMask, expr: ExprRef) -> VortexResult<RowMask> {
@@ -128,7 +127,7 @@ mod tests {
             .map(IntoArray::into_array)],
         )
         .unwrap();
-        (ScanExecutor::inline(Arc::new(segments)), layout)
+        (ScanExecutor::new(Arc::new(segments)), layout)
     }
 
     #[test]

vortex-layout/src/scan/executor/mod.rs

@@ -0,0 +1,49 @@
+#[cfg(feature = "tokio")]
+mod tokio;
+
+mod threads;
+
+use std::future::Future;
+
+use futures::future::BoxFuture;
+pub use threads::*;
+#[cfg(feature = "tokio")]
+pub use tokio::*;
+use vortex_error::VortexResult;
+
+pub trait Executor {
+    /// Spawns a future to run on a different runtime. The shouldn't be polled to make progress.
+    fn spawn<F>(&self, f: F) -> VortexResult<BoxFuture<'static, VortexResult<F::Output>>>
+    where
+        F: Future + Send + 'static,
+        <F as Future>::Output: Send + 'static;
+}
+
+/// Generic wrapper around different async runtimes. Used to spawn async tasks to run in the background, concurrently with other tasks.
+#[derive(Clone)]
+pub enum TaskExecutor {
+    Threads(ThreadsExecutor),
+    #[cfg(feature = "tokio")]
+    Tokio(TokioExecutor),
+}
+
+pub enum TaskExecutorMode {
+    Threads,
+    #[cfg(feature = "tokio")]
+    Tokio,
+}
+
+#[async_trait::async_trait]
+impl Executor for TaskExecutor {
+    fn spawn<F>(&self, f: F) -> VortexResult<BoxFuture<'static, VortexResult<F::Output>>>
+    where
+        F: Future + Send + 'static,
+        <F as Future>::Output: Send + 'static,
+    {
+        match self {
+            TaskExecutor::Threads(threads_executor) => threads_executor.spawn(f),
+            #[cfg(feature = "tokio")]
+            TaskExecutor::Tokio(tokio_executor) => tokio_executor.spawn(f),
+        }
+    }
+}

vortex-layout/src/scan/executor/threads.rs

@@ -0,0 +1,96 @@
+use std::future::Future;
+use std::num::NonZeroUsize;
+use std::sync::Arc;
+
+use futures::channel::oneshot;
+use futures::future::BoxFuture;
+use futures::{FutureExt as _, TryFutureExt as _};
+use vortex_error::{vortex_err, VortexResult};
+
+use super::Executor;
+
+trait Task {
+    fn run(self: Box<Self>);
+}
+
+struct ExecutorTask<F, R> {
+    task: F,
+    result: oneshot::Sender<R>,
+}
+
+impl<F, R> Task for ExecutorTask<F, R>
+where
+    F: Future<Output = R> + Send,
+    R: Send,
+{
+    fn run(self: Box<Self>) {
+        let Self { task, result } = *self;
+        futures::executor::block_on(async move {
+            let output = task.await;
+            _ = result.send(output);
+        })
+    }
+}
+
+/// Multithreaded task executor, runs tasks on a dedicated thread pool.
+#[derive(Clone, Default)]
+pub struct ThreadsExecutor {
+    inner: Arc<Inner>,
+}
+
+impl ThreadsExecutor {
+    pub fn new(num_threads: NonZeroUsize) -> Self {
+        Self {
+            inner: Arc::new(Inner::new(num_threads)),
+        }
+    }
+}
+
+struct Inner {
+    submitter: flume::Sender<Box<dyn Task + Send>>,
+}
+
+impl Default for Inner {
+    fn default() -> Self {
+        // Safety:
+        // 1 isn't 0
+        Self::new(unsafe { NonZeroUsize::new_unchecked(1) })
+    }
+}
+
+impl Inner {
+    fn new(num_threads: NonZeroUsize) -> Self {
+        let (tx, rx) = flume::unbounded::<Box<dyn Task + Send>>();
+        (0..num_threads.get()).for_each(|_| {
+            let rx = rx.clone();
+            std::thread::spawn(move || {
+                // The channel errors if all senders are dropped, which means we probably don't care about the task anymore,
+                // and we can break and let the thread end.
+                while let Ok(task) = rx.recv() {
+                    task.run()
+                }
+            });
+        });
+
+        Self { submitter: tx }
+    }
+}
+
+impl Executor for ThreadsExecutor {
+    fn spawn<F>(&self, f: F) -> VortexResult<BoxFuture<'static, VortexResult<F::Output>>>
+    where
+        F: Future + Send + 'static,
+        <F as Future>::Output: Send + 'static,
+    {
+        let (tx, rx) = oneshot::channel();
+        let task = Box::new(ExecutorTask {
+            task: f,
+            result: tx,
+        });
+        self.inner
+            .submitter
+            .send(task)
+            .map_err(|e| vortex_err!("Failed to submit work: {e}"))?;
+        Ok(rx.map_err(|e| vortex_err!("Future canceled: {e}")).boxed())
+    }
+}

vortex-layout/src/scan/executor/tokio.rs

@@ -0,0 +1,29 @@
+use std::future::Future;
+
+use futures::future::BoxFuture;
+use futures::{FutureExt, TryFutureExt};
+use tokio::runtime::Handle;
+use vortex_error::{VortexError, VortexResult};
+
+use super::Executor;
+
+/// Tokio-based async task executor, runs task on the provided runtime.
+#[derive(Clone)]
+pub struct TokioExecutor(Handle);
+
+impl TokioExecutor {
+    pub fn new(handle: Handle) -> Self {
+        Self(handle)
+    }
+}
+
+#[async_trait::async_trait]
+impl Executor for TokioExecutor {
+    fn spawn<F>(&self, f: F) -> VortexResult<BoxFuture<'static, VortexResult<F::Output>>>
+    where
+        F: Future + Send + 'static,
+        <F as Future>::Output: Send + 'static,
+    {
+        Ok(self.0.spawn(f).map_err(VortexError::from).boxed())
+    }
+}