compute pool

src/daft-io/src/lib.rs

@@ -465,6 +465,25 @@ impl Runtime {
         Arc::new(Self { runtime })
     }
 
+    async fn execute_io_task<F>(future: F) -> DaftResult<F::Output>
+    where
+        F: Future + Send + 'static,
+        F::Output: Send + 'static,
+    {
+        AssertUnwindSafe(future).catch_unwind().await.map_err(|e| {
+            let s = if let Some(s) = e.downcast_ref::<String>() {
+                s.clone()
+            } else if let Some(s) = e.downcast_ref::<&str>() {
+                s.to_string()
+            } else {
+                "unknown internal error".to_string()
+            };
+            DaftError::ComputeError(format!(
+                "Caught panic when spawning blocking task in io pool {s})"
+            ))
+        })
+    }
+
     /// Similar to tokio's Runtime::block_on but requires static lifetime + Send
     /// You should use this when you are spawning IO tasks from an Expression Evaluator or in the Executor
     pub fn block_on_io_pool<F>(&self, future: F) -> DaftResult<F::Output>
@@ -474,26 +493,30 @@ impl Runtime {
     {
         let (tx, rx) = oneshot::channel();
         let _join_handle = self.spawn(async move {
-            let task_output = AssertUnwindSafe(future).catch_unwind().await.map_err(|e| {
-                let s = if let Some(s) = e.downcast_ref::<String>() {
-                    s.clone()
-                } else if let Some(s) = e.downcast_ref::<&str>() {
-                    s.to_string()
-                } else {
-                    "unknown internal error".to_string()
-                };
-                DaftError::ComputeError(format!(
-                    "Caught panic when spawning blocking task in io pool {s})"
-                ))
-            });
-
+            let task_output = Self::execute_io_task(future).await;
             if tx.send(task_output).is_err() {
-                log::warn!("Spawned task output ignored: receiver dropped")
+                log::warn!("Spawned task output ignored: receiver dropped");
             }
         });
         rx.recv().expect("Spawned task transmitter dropped")
     }
 
+    /// Similar to block_on_io_pool, but is async and can be awaited
+    pub async fn block_on_io_pool_async<F>(&self, future: F) -> DaftResult<F::Output>
+    where
+        F: Future + Send + 'static,
+        F::Output: Send + 'static,
+    {
+        let (tx, rx) = oneshot::channel();
+        let _join_handle = self.spawn(async move {
+            let task_output = Self::execute_io_task(future).await;
+            if tx.send(task_output).is_err() {
+                log::warn!("Spawned task output ignored: receiver dropped");
+            }
+        });
+        rx.await.expect("Spawned task transmitter dropped")
+    }
+
     /// Blocks current thread to compute future. Can not be called in tokio runtime context
     ///
     pub fn block_on_current_thread<F: Future>(&self, future: F) -> F::Output {

src/daft-local-execution/src/intermediate_ops/aggregate.rs

@@ -29,7 +29,7 @@ impl IntermediateOperator for AggregateOperator {
         &self,
         _idx: usize,
         input: &PipelineResultType,
-        _state: Option<&mut Box<dyn IntermediateOperatorState>>,
+        _state: &mut dyn IntermediateOperatorState,
     ) -> DaftResult<IntermediateOperatorResult> {
         let out = input.as_data().agg(&self.agg_exprs, &self.group_by)?;
         Ok(IntermediateOperatorResult::NeedMoreInput(Some(Arc::new(

src/daft-local-execution/src/intermediate_ops/anti_semi_hash_join_probe.rs

@@ -100,12 +100,11 @@ impl IntermediateOperator for AntiSemiProbeOperator {
         &self,
         idx: usize,
         input: &PipelineResultType,
-        state: Option<&mut Box<dyn IntermediateOperatorState>>,
+        state: &mut dyn IntermediateOperatorState,
     ) -> DaftResult<IntermediateOperatorResult> {
         match idx {
             0 => {
                 let state = state
-                    .expect("AntiSemiProbeOperator should have state")
                     .as_any_mut()
                     .downcast_mut::<AntiSemiProbeState>()
                     .expect("AntiSemiProbeOperator state should be AntiSemiProbeState");
@@ -115,7 +114,6 @@ impl IntermediateOperator for AntiSemiProbeOperator {
             }
             _ => {
                 let state = state
-                    .expect("AntiSemiProbeOperator should have state")
                     .as_any_mut()
                     .downcast_mut::<AntiSemiProbeState>()
                     .expect("AntiSemiProbeOperator state should be AntiSemiProbeState");
@@ -133,7 +131,7 @@ impl IntermediateOperator for AntiSemiProbeOperator {
         "AntiSemiProbeOperator"
     }
 
-    fn make_state(&self) -> Option<Box<dyn IntermediateOperatorState>> {
-        Some(Box::new(AntiSemiProbeState::Building))
+    fn make_state(&self) -> Box<dyn IntermediateOperatorState> {
+        Box::new(AntiSemiProbeState::Building)
     }
 }

src/daft-local-execution/src/intermediate_ops/filter.rs

@@ -25,7 +25,7 @@ impl IntermediateOperator for FilterOperator {
         &self,
         _idx: usize,
         input: &PipelineResultType,
-        _state: Option<&mut Box<dyn IntermediateOperatorState>>,
+        _state: &mut dyn IntermediateOperatorState,
     ) -> DaftResult<IntermediateOperatorResult> {
         let out = input.as_data().filter(&[self.predicate.clone()])?;
         Ok(IntermediateOperatorResult::NeedMoreInput(Some(Arc::new(

src/daft-local-execution/src/intermediate_ops/hash_join_probe.rs

@@ -231,12 +231,11 @@ impl IntermediateOperator for HashJoinProbeOperator {
         &self,
         idx: usize,
         input: &PipelineResultType,
-        state: Option<&mut Box<dyn IntermediateOperatorState>>,
+        state: &mut dyn IntermediateOperatorState,
     ) -> DaftResult<IntermediateOperatorResult> {
         match idx {
             0 => {
                 let state = state
-                    .expect("HashJoinProbeOperator should have state")
                     .as_any_mut()
                     .downcast_mut::<HashJoinProbeState>()
                     .expect("HashJoinProbeOperator state should be HashJoinProbeState");
@@ -246,7 +245,6 @@ impl IntermediateOperator for HashJoinProbeOperator {
             }
             _ => {
                 let state = state
-                    .expect("HashJoinProbeOperator should have state")
                     .as_any_mut()
                     .downcast_mut::<HashJoinProbeState>()
                     .expect("HashJoinProbeOperator state should be HashJoinProbeState");
@@ -267,7 +265,7 @@ impl IntermediateOperator for HashJoinProbeOperator {
         "HashJoinProbeOperator"
     }
 
-    fn make_state(&self) -> Option<Box<dyn IntermediateOperatorState>> {
-        Some(Box::new(HashJoinProbeState::Building))
+    fn make_state(&self) -> Box<dyn IntermediateOperatorState> {
+        Box::new(HashJoinProbeState::Building)
     }
 }

src/daft-local-execution/src/intermediate_ops/intermediate_op.rs

@@ -3,6 +3,7 @@ use std::sync::Arc;
 use common_display::tree::TreeDisplay;
 use common_error::DaftResult;
 use daft_micropartition::MicroPartition;
+use tokio::sync::Mutex;
 use tracing::{info_span, instrument};
 
 use super::buffer::OperatorBuffer;
@@ -16,6 +17,12 @@ use crate::{
 pub trait IntermediateOperatorState: Send + Sync {
     fn as_any_mut(&mut self) -> &mut dyn std::any::Any;
 }
+pub struct DefaultIntermediateOperatorState {}
+impl IntermediateOperatorState for DefaultIntermediateOperatorState {
+    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
+        self
+    }
+}
 
 pub enum IntermediateOperatorResult {
     NeedMoreInput(Option<Arc<MicroPartition>>),
@@ -28,11 +35,11 @@ pub trait IntermediateOperator: Send + Sync {
         &self,
         idx: usize,
         input: &PipelineResultType,
-        state: Option<&mut Box<dyn IntermediateOperatorState>>,
+        state: &mut dyn IntermediateOperatorState,
     ) -> DaftResult<IntermediateOperatorResult>;
     fn name(&self) -> &'static str;
-    fn make_state(&self) -> Option<Box<dyn IntermediateOperatorState>> {
-        None
+    fn make_state(&self) -> Box<dyn IntermediateOperatorState> {
+        Box::new(DefaultIntermediateOperatorState {})
     }
 }
 
@@ -75,11 +82,20 @@ impl IntermediateNode {
         rt_context: Arc<RuntimeStatsContext>,
     ) -> DaftResult<()> {
         let span = info_span!("IntermediateOp::execute");
-        let mut state = op.make_state();
+        let compute_runtime = crate::get_compute_runtime()?;
+        let state = Arc::new(Mutex::new(op.make_state()));
         while let Some((idx, morsel)) = receiver.recv().await {
             loop {
-                let result =
-                    rt_context.in_span(&span, || op.execute(idx, &morsel, state.as_mut()))?;
+                let state = state.clone();
+                let op = op.clone();
+                let morsel = morsel.clone();
+                let span = span.clone();
+                let rt_context = rt_context.clone();
+                let fut = async move {
+                    let mut state_guard = state.lock().await;
+                    rt_context.in_span(&span, || op.execute(idx, &morsel, state_guard.as_mut()))
+                };
+                let result = compute_runtime.block_on_compute_pool(fut).await??;
                 match result {
                     IntermediateOperatorResult::NeedMoreInput(Some(mp)) => {
                         let _ = sender.send(mp.into()).await;

src/daft-local-execution/src/intermediate_ops/project.rs

@@ -25,7 +25,7 @@ impl IntermediateOperator for ProjectOperator {
         &self,
         _idx: usize,
         input: &PipelineResultType,
-        _state: Option<&mut Box<dyn IntermediateOperatorState>>,
+        _state: &mut dyn IntermediateOperatorState,
     ) -> DaftResult<IntermediateOperatorResult> {
         let out = input.as_data().eval_expression_list(&self.projection)?;
         Ok(IntermediateOperatorResult::NeedMoreInput(Some(Arc::new(

src/daft-local-execution/src/lib.rs

@@ -6,14 +6,99 @@ mod run;
 mod runtime_stats;
 mod sinks;
 mod sources;
+use std::{
+    future::Future,
+    panic::AssertUnwindSafe,
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc, OnceLock,
+    },
+};
+
 use common_error::{DaftError, DaftResult};
+use futures::FutureExt;
 use lazy_static::lazy_static;
 pub use run::NativeExecutor;
 use snafu::{futures::TryFutureExt, Snafu};
 lazy_static! {
     pub static ref NUM_CPUS: usize = std::thread::available_parallelism().unwrap().get();
 }
 
+static COMPUTE_RUNTIME: OnceLock<RuntimeRef> = OnceLock::new();
+
+pub type RuntimeRef = Arc<Runtime>;
+
+pub struct Runtime {
+    runtime: tokio::runtime::Runtime,
+}
+
+impl Runtime {
+    fn new(runtime: tokio::runtime::Runtime) -> RuntimeRef {
+        Arc::new(Self { runtime })
+    }
+
+    pub async fn block_on_compute_pool<F>(&self, future: F) -> DaftResult<F::Output>
+    where
+        F: Future + Send + 'static,
+        F::Output: Send + 'static,
+    {
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        let _join_handle = self.spawn(async move {
+            let task_output = AssertUnwindSafe(future).catch_unwind().await.map_err(|e| {
+                let s = if let Some(s) = e.downcast_ref::<String>() {
+                    s.clone()
+                } else if let Some(s) = e.downcast_ref::<&str>() {
+                    s.to_string()
+                } else {
+                    "unknown internal error".to_string()
+                };
+                DaftError::ComputeError(format!(
+                    "Caught panic when spawning blocking task in compute pool {s})"
+                ))
+            });
+
+            if tx.send(task_output).is_err() {
+                log::warn!("Spawned task output ignored: receiver dropped")
+            }
+        });
+        rx.await.expect("Compute pool receiver dropped")
+    }
+
+    pub fn spawn<F>(&self, future: F) -> JoinHandle<F::Output>
+    where
+        F: Future + Send + 'static,
+        F::Output: Send + 'static,
+    {
+        self.runtime.spawn(future)
+    }
+}
+
+fn init_runtime(num_threads: usize) -> Arc<Runtime> {
+    std::thread::spawn(move || {
+        Runtime::new(
+            tokio::runtime::Builder::new_multi_thread()
+                .worker_threads(num_threads)
+                .enable_all()
+                .thread_name_fn(|| {
+                    static ATOMIC_ID: AtomicUsize = AtomicUsize::new(0);
+                    let id = ATOMIC_ID.fetch_add(1, Ordering::SeqCst);
+                    format!("Executor-Worker-{}", id)
+                })
+                .build()
+                .unwrap(),
+        )
+    })
+    .join()
+    .unwrap()
+}
+
+pub fn get_compute_runtime() -> DaftResult<RuntimeRef> {
+    let runtime = COMPUTE_RUNTIME
+        .get_or_init(|| init_runtime(*NUM_CPUS))
+        .clone();
+    Ok(runtime)
+}
+
 pub struct ExecutionRuntimeHandle {
     worker_set: tokio::task::JoinSet<crate::Result<()>>,
     default_morsel_size: usize,
@@ -51,6 +136,7 @@ impl ExecutionRuntimeHandle {
 
 #[cfg(feature = "python")]
 use pyo3::prelude::*;
+use tokio::task::JoinHandle;
 
 #[derive(Debug, Snafu)]
 pub enum Error {

src/daft-local-execution/src/run.rs

@@ -2,10 +2,7 @@ use std::{
     collections::HashMap,
     fs::File,
     io::Write,
-    sync::{
-        atomic::{AtomicUsize, Ordering},
-        Arc,
-    },
+    sync::Arc,
     time::{SystemTime, UNIX_EPOCH},
 };
 
@@ -124,14 +121,8 @@ pub fn run_local(
     let mut pipeline = physical_plan_to_pipeline(physical_plan, &psets)?;
     let (tx, rx) = create_channel(results_buffer_size.unwrap_or(1));
     let handle = std::thread::spawn(move || {
-        let runtime = tokio::runtime::Builder::new_multi_thread()
+        let runtime = tokio::runtime::Builder::new_current_thread()
             .enable_all()
-            .max_blocking_threads(10)
-            .thread_name_fn(|| {
-                static ATOMIC_ID: AtomicUsize = AtomicUsize::new(0);
-                let id = ATOMIC_ID.fetch_add(1, Ordering::SeqCst);
-                format!("Executor-Worker-{}", id)
-            })
             .build()
             .expect("Failed to create tokio runtime");
         runtime.block_on(async {