Add execution concurrency == 2 * io concurrency (#2023)

Author: gatesn

vortex-file/src/exec/inline.rs

@@ -1,21 +1,32 @@
+use std::future::ready;
+
 use futures_util::future::BoxFuture;
 use futures_util::stream::BoxStream;
-use futures_util::{FutureExt, StreamExt};
+use futures_util::StreamExt;
 use vortex_array::ArrayData;
 use vortex_error::VortexResult;
 
 use crate::exec::ExecDriver;
 
 /// An [`ExecDriver`] implementation that awaits the futures inline using the caller's runtime.
-pub struct InlineDriver;
+pub struct InlineDriver {
+    concurrency: usize,
+}
+
+impl InlineDriver {
+    pub fn with_concurrency(concurrency: usize) -> Self {
+        Self { concurrency }
+    }
+}
 
 impl ExecDriver for InlineDriver {
     fn drive(
         &self,
         stream: BoxStream<'static, BoxFuture<'static, VortexResult<Option<ArrayData>>>>,
     ) -> BoxStream<'static, VortexResult<ArrayData>> {
         stream
-            .filter_map(|future| future.map(|result| result.transpose()))
+            .buffered(self.concurrency)
+            .filter_map(|result| ready(result.transpose()))
             .boxed()
     }
 }

vortex-file/src/exec/tokio.rs

@@ -10,23 +10,21 @@ use vortex_error::{vortex_err, VortexResult};
 use crate::exec::ExecDriver;
 
 /// An [`ExecDriver`] implementation that spawns the futures onto a Tokio runtime.
-pub struct TokioDriver(pub Handle);
+pub struct TokioDriver {
+    pub handle: Handle,
+    pub concurrency: usize,
+}
 
 impl ExecDriver for TokioDriver {
     fn drive(
         &self,
         stream: BoxStream<'static, BoxFuture<'static, VortexResult<Option<ArrayData>>>>,
     ) -> BoxStream<'static, VortexResult<ArrayData>> {
-        let handle = self.0.clone();
-
-        // This is how many file splits to make progress on at once. While I/O is resolving for
-        // the first, we may as well find out the segments required by the next.
-        // TODO(ngates): I picked this number somewhat arbitrarily :)
-        let concurrency = 2 * handle.metrics().num_workers();
+        let handle = self.handle.clone();
 
         stream
             .map(move |future| handle.spawn(future))
-            .buffered(concurrency)
+            .buffered(self.concurrency)
             .map(|result| match result {
                 Ok(result) => result,
                 Err(e) => Err(vortex_err!("Failed to join Tokio result {}", e)),

vortex-file/src/open/exec.rs

@@ -1,8 +1,5 @@
 use std::sync::Arc;
 
-#[cfg(feature = "tokio")]
-use tokio::runtime::Handle;
-
 use crate::exec::inline::InlineDriver;
 #[cfg(feature = "tokio")]
 use crate::exec::tokio::TokioDriver;
@@ -19,29 +16,35 @@ pub enum ExecutionMode {
     RayonThreadPool(Arc<rayon::ThreadPool>),
     /// Spawns the tasks onto a provided Tokio runtime.
     #[cfg(feature = "tokio")]
-    TokioRuntime(Handle),
+    TokioRuntime(tokio::runtime::Handle),
+}
+
+#[allow(clippy::derivable_impls)]
+impl Default for ExecutionMode {
+    fn default() -> Self {
+        // Default to tokio-specific behavior if its enabled and there's a runtime running.
+        #[cfg(feature = "tokio")]
+        if let Ok(h) = tokio::runtime::Handle::try_current() {
+            return ExecutionMode::TokioRuntime(h);
+        }
+
+        ExecutionMode::Inline
+    }
 }
 
 impl ExecutionMode {
-    pub fn into_driver(self) -> Arc<dyn ExecDriver> {
+    pub fn into_driver(self, concurrency: usize) -> Arc<dyn ExecDriver> {
         match self {
-            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)
-            }
+            ExecutionMode::Inline => Arc::new(InlineDriver::with_concurrency(concurrency)),
             #[cfg(feature = "rayon")]
             ExecutionMode::RayonThreadPool(_) => {
                 todo!()
             }
             #[cfg(feature = "tokio")]
-            ExecutionMode::TokioRuntime(handle) => Arc::new(TokioDriver(handle)),
+            ExecutionMode::TokioRuntime(handle) => Arc::new(TokioDriver {
+                handle,
+                concurrency,
+            }),
         }
     }
 }

vortex-file/src/open/mod.rs

@@ -42,6 +42,7 @@ pub struct VortexOpenOptions {
     execution_mode: Option<ExecutionMode>,
     // TODO(ngates): allow fully configurable I/O driver.
     io_concurrency: usize,
+    exec_concurrency: Option<usize>,
 }
 
 impl VortexOpenOptions {
@@ -56,7 +57,8 @@ impl VortexOpenOptions {
             segment_cache: None,
             execution_mode: None,
             // TODO(ngates): pick some numbers...
-            io_concurrency: 16,
+            io_concurrency: 10,
+            exec_concurrency: None,
         }
     }
 
@@ -111,6 +113,21 @@ impl VortexOpenOptions {
         self.execution_mode = Some(execution_mode);
         self
     }
+
+    /// Configure the number of concurrent I/O requests.
+    pub fn with_io_concurrency(mut self, io_concurrency: usize) -> Self {
+        self.io_concurrency = io_concurrency;
+        self
+    }
+
+    /// Override the default split-by concurrency.
+    ///
+    /// It is recommended to use more split-by concurrency than I/O concurrency to ensure there
+    /// are always I/O operations enqueued.
+    pub fn with_exec_concurrency(mut self, exec_concurrency: usize) -> Self {
+        self.exec_concurrency = Some(exec_concurrency);
+        self
+    }
 }
 
 impl VortexOpenOptions {
@@ -143,8 +160,8 @@ impl VortexOpenOptions {
         // Set up the execution driver.
         let exec_driver = self
             .execution_mode
-            .unwrap_or(ExecutionMode::Inline)
-            .into_driver();
+            .unwrap_or_default()
+            .into_driver(self.exec_concurrency.unwrap_or(self.io_concurrency * 2));
 
         // Finally, create the VortexFile.
         Ok(VortexFile {