feature flag ExecDriver implementations (#2005)

1. Feature flag tokio/rayon based `ExecDriver` implementation
2. If tokio is available AND running within the context of a runtime, use it.

Author: AdamGS

vortex-datafusion/Cargo.toml

@@ -40,7 +40,7 @@ vortex-buffer = { workspace = true }
 vortex-dtype = { workspace = true }
 vortex-error = { workspace = true, features = ["datafusion"] }
 vortex-expr = { workspace = true, features = ["datafusion"] }
-vortex-file = { workspace = true, features = ["object_store"] }
+vortex-file = { workspace = true, features = ["object_store", "tokio"] }
 vortex-io = { workspace = true, features = ["object_store", "tokio"] }
 vortex-scan = { workspace = true }
 

vortex-file/Cargo.toml

@@ -27,8 +27,8 @@ futures-util = { workspace = true }
 itertools = { workspace = true }
 log = { workspace = true }
 pin-project-lite = { workspace = true }
-rayon = { workspace = true }
-tokio = { workspace = true, features = ["rt"] }
+rayon = { workspace = true, optional = true }
+tokio = { workspace = true, features = ["rt"], optional = true }
 tracing = { workspace = true, optional = true }
 vortex-array = { workspace = true }
 vortex-buffer = { workspace = true }

vortex-file/src/exec/mod.rs

@@ -1,4 +1,6 @@
 pub mod inline;
+
+#[cfg(feature = "tokio")]
 pub mod tokio;
 
 use futures_util::future::BoxFuture;

vortex-file/src/open/exec.rs

@@ -1,6 +1,10 @@
 use std::sync::Arc;
 
+#[cfg(feature = "tokio")]
+use tokio::runtime::Handle;
+
 use crate::exec::inline::InlineDriver;
+#[cfg(feature = "tokio")]
 use crate::exec::tokio::TokioDriver;
 use crate::exec::ExecDriver;
 
@@ -11,20 +15,32 @@ pub enum ExecutionMode {
     /// [`vortex_array::stream::ArrayStream`]. In other words, uses the same runtime.
     Inline,
     /// Spawns the tasks onto a provided Rayon thread pool.
-    // TODO(ngates): feature-flag this dependency.
+    #[cfg(feature = "rayon")]
     RayonThreadPool(Arc<rayon::ThreadPool>),
     /// Spawns the tasks onto a provided Tokio runtime.
-    // TODO(ngates): feature-flag this dependency.
-    TokioRuntime(tokio::runtime::Handle),
+    #[cfg(feature = "tokio")]
+    TokioRuntime(Handle),
 }
 
 impl ExecutionMode {
     pub fn into_driver(self) -> Arc<dyn ExecDriver> {
         match self {
-            ExecutionMode::Inline => Arc::new(InlineDriver),
+            ExecutionMode::Inline => {
+                // Default to tokio-specific behavior if its enabled and there's a runtime running.
+                #[cfg(feature = "tokio")]
+                match Handle::try_current() {
+                    Ok(h) => Arc::new(TokioDriver(h)),
+                    Err(_) => Arc::new(InlineDriver),
+                }
+
+                #[cfg(not(feature = "tokio"))]
+                Arc::new(InlineDriver)
+            }
+            #[cfg(feature = "rayon")]
             ExecutionMode::RayonThreadPool(_) => {
                 todo!()
             }
+            #[cfg(feature = "tokio")]
             ExecutionMode::TokioRuntime(handle) => Arc::new(TokioDriver(handle)),
         }
     }

vortex-file/src/read/record_batch_reader.rs

@@ -24,6 +24,7 @@ pub trait AsyncRuntime {
     fn block_on<F: Future>(&self, fut: F) -> F::Output;
 }
 
+#[cfg(feature = "tokio")]
 impl AsyncRuntime for tokio::runtime::Runtime {
     fn block_on<F: Future>(&self, fut: F) -> F::Output {
         self.block_on(fut)

vortex/Cargo.toml

@@ -50,7 +50,7 @@ vortex-zigzag = { workspace = true }
 vortex-roaring = { workspace = true }
 
 [features]
-tokio = ["vortex-io/tokio"]
+tokio = ["vortex-io/tokio", "vortex-file/tokio"]
 object_store = ["vortex-file/object_store"]
 parquet = ["vortex-error/parquet"]
 python = ["vortex-error/python"]