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::{Executor, 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 = Executor::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_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

@@ -212,6 +212,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/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> {

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> {

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

@@ -0,0 +1,43 @@
+#[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 Spawn {
+    // Spawns a future to run on a different runtime. The returning future should be polled to ensure its running.
+    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. Can be used to spawn futures to run in the background, concurrently with other tasks.
+#[derive(Clone)]
+pub enum Executor {
+    Threads(ThreadsExecutor),
+    #[cfg(feature = "tokio")]
+    Tokio(TokioExecutor),
+}
+
+#[async_trait::async_trait]
+impl Spawn for Executor {
+    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 {
+            Executor::Threads(threads_executor) => threads_executor.spawn(f),
+            #[cfg(feature = "tokio")]
+            Executor::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::Spawn;
+
+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);
+        })
+    }
+}
+
+#[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 || loop {
+                match rx.recv() {
+                    Ok(task) => task.run(),
+                    // we error if all senders dropped, which means we probably don't care about the task anymore,
+                    // and we can break and let the thread end.
+                    Err(_e) => break,
+                }
+            });
+        });
+
+        Self { submitter: tx }
+    }
+}
+
+impl Spawn 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,28 @@
+use std::future::Future;
+
+use futures::future::BoxFuture;
+use futures::{FutureExt, TryFutureExt};
+use tokio::runtime::Handle;
+use vortex_error::{VortexError, VortexResult};
+
+use super::Spawn;
+
+#[derive(Clone)]
+pub struct TokioExecutor(Handle);
+
+impl TokioExecutor {
+    pub fn new(handle: Handle) -> Self {
+        Self(handle)
+    }
+}
+
+#[async_trait::async_trait]
+impl Spawn 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())
+    }
+}