par eval expr

Author: EC2 Default User

Cargo.lock

@@ -4,6 +4,7 @@ futures = {workspace = true}
 lazy_static = {workspace = true}
 log = {workspace = true}
 oneshot = "0.1.8"
+thread-priority = "1.2.0"
 tokio = {workspace = true}
 
 [lints]

src/common/runtime/Cargo.toml

@@ -12,6 +12,7 @@ use std::{
 use common_error::{DaftError, DaftResult};
 use futures::FutureExt;
 use lazy_static::lazy_static;
+use thread_priority::set_current_thread_priority;
 use tokio::{
     runtime::{Handle, RuntimeFlavor},
     task::JoinSet,
@@ -149,6 +150,9 @@ fn init_compute_runtime() -> RuntimeRef {
                 let id = COMPUTE_THREAD_ATOMIC_ID.fetch_add(1, Ordering::SeqCst);
                 format!("Compute-Thread-{}", id)
             })
+            .on_thread_start(|| {
+                set_current_thread_priority(thread_priority::ThreadPriority::Min).unwrap();
+            })
             .max_blocking_threads(*COMPUTE_RUNTIME_MAX_BLOCKING_THREADS);
         Runtime::new(builder.build().unwrap(), PoolType::Compute)
     })

src/common/runtime/src/lib.rs

@@ -1219,6 +1219,38 @@ impl Expr {
             _ => None,
         }
     }
+
+    pub fn has_compute(&self) -> bool {
+        match self {
+            Self::Column(..) => false,
+            Self::Literal(..) => false,
+            Self::Subquery(..) => false,
+            Self::Exists(..) => false,
+            Self::OuterReferenceColumn(..) => false,
+            Self::Function { .. } => true,
+            Self::ScalarFunction(..) => true,
+            Self::Agg(_) => true,
+            Self::Alias(expr, ..) => expr.has_compute(),
+            Self::Cast(expr, ..) => expr.has_compute(),
+            Self::Not(expr) => expr.has_compute(),
+            Self::IsNull(expr) => expr.has_compute(),
+            Self::NotNull(expr) => expr.has_compute(),
+            Self::FillNull(expr, fill_value) => expr.has_compute() || fill_value.has_compute(),
+            Self::IsIn(expr, items) => {
+                expr.has_compute() || items.iter().any(|item| item.has_compute())
+            }
+            Self::Between(expr, lower, upper) => {
+                expr.has_compute() || lower.has_compute() || upper.has_compute()
+            }
+            Self::BinaryOp { left, right, .. } => left.has_compute() || right.has_compute(),
+            Self::IfElse {
+                if_true,
+                if_false,
+                predicate,
+            } => if_true.has_compute() || if_false.has_compute() || predicate.has_compute(),
+            Self::InSubquery(expr, _) => expr.has_compute(),
+        }
+    }
 }
 
 #[derive(Display, Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize, Hash)]

src/daft-dsl/src/expr/mod.rs

@@ -23,14 +23,15 @@ daft-table = {path = "../daft-table", default-features = false}
 daft-writers = {path = "../daft-writers", default-features = false}
 futures = {workspace = true}
 indexmap = {workspace = true}
+kanal = "0.1.0-pre8"
 lazy_static = {workspace = true}
 log = {workspace = true}
-loole = "0.4.0"
 num-format = "0.4.4"
 pin-project = "1"
 pyo3 = {workspace = true, optional = true}
 snafu = {workspace = true}
 tokio = {workspace = true}
+tokio-stream = {workspace = true}
 tokio-util = {workspace = true}
 tracing = {workspace = true}
 

src/daft-local-execution/Cargo.toml

@@ -1,29 +1,25 @@
 #[derive(Clone)]
-pub(crate) struct Sender<T>(loole::Sender<T>);
+pub(crate) struct Sender<T>(kanal::AsyncSender<T>);
 impl<T> Sender<T> {
-    pub(crate) async fn send(&self, val: T) -> Result<(), loole::SendError<T>> {
-        self.0.send_async(val).await
+    pub(crate) async fn send(&self, val: T) -> Result<(), kanal::SendError> {
+        self.0.send(val).await
     }
 }
 
 #[derive(Clone)]
-pub(crate) struct Receiver<T>(loole::Receiver<T>);
+pub(crate) struct Receiver<T>(kanal::AsyncReceiver<T>);
 impl<T> Receiver<T> {
     pub(crate) async fn recv(&self) -> Option<T> {
-        self.0.recv_async().await.ok()
+        self.0.recv().await.ok()
     }
 
-    pub(crate) fn blocking_recv(&self) -> Option<T> {
-        self.0.recv().ok()
-    }
-
-    pub(crate) fn into_inner(self) -> loole::Receiver<T> {
+    pub(crate) fn into_inner(self) -> kanal::AsyncReceiver<T> {
         self.0
     }
 }
 
 pub(crate) fn create_channel<T: Clone>(buffer_size: usize) -> (Sender<T>, Receiver<T>) {
-    let (tx, rx) = loole::bounded(buffer_size);
+    let (tx, rx) = kanal::bounded_async::<T>(buffer_size);
     (Sender(tx), Receiver(rx))
 }
 
@@ -36,7 +32,7 @@ pub(crate) fn create_ordering_aware_receiver_channel<T: Clone>(
 ) -> (Vec<Sender<T>>, OrderingAwareReceiver<T>) {
     match ordered {
         true => {
-            let (senders, receiver) = (0..buffer_size).map(|_| create_channel::<T>(1)).unzip();
+            let (senders, receiver) = (0..buffer_size).map(|_| create_channel::<T>(0)).unzip();
             (
                 senders,
                 OrderingAwareReceiver::InOrder(RoundRobinReceiver::new(receiver)),

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

@@ -95,7 +95,7 @@ impl DispatchSpawner for RoundRobinDispatcher {
         runtime_handle: &mut RuntimeHandle,
     ) -> SpawnedDispatchResult {
         let (worker_senders, worker_receivers): (Vec<_>, Vec<_>) =
-            (0..num_workers).map(|_| create_channel(1)).unzip();
+            (0..num_workers).map(|_| create_channel(0)).unzip();
         let morsel_size = self.morsel_size;
         let task = runtime_handle.spawn(async move {
             Self::dispatch_inner(worker_senders, input_receivers, morsel_size).await
@@ -213,7 +213,7 @@ impl DispatchSpawner for PartitionedDispatcher {
         runtime_handle: &mut RuntimeHandle,
     ) -> SpawnedDispatchResult {
         let (worker_senders, worker_receivers): (Vec<_>, Vec<_>) =
-            (0..num_workers).map(|_| create_channel(1)).unzip();
+            (0..num_workers).map(|_| create_channel(0)).unzip();
         let partition_by = self.partition_by.clone();
         let dispatch_task = runtime_handle.spawn(async move {
             Self::dispatch_inner(worker_senders, input_receivers, partition_by).await

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

@@ -212,7 +212,7 @@ impl PipelineNode for IntermediateNode {
         let num_workers = op.max_concurrency().context(PipelineExecutionSnafu {
             node_name: self.name(),
         })?;
-        let (destination_sender, destination_receiver) = create_channel(1);
+        let (destination_sender, destination_receiver) = create_channel(0);
         let counting_sender = CountingSender::new(destination_sender, self.runtime_stats.clone());
 
         let dispatch_spawner = self

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

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::{cmp::max, sync::Arc};
 
 use common_error::{DaftError, DaftResult};
 use common_runtime::RuntimeRef;
@@ -12,15 +12,73 @@ use super::intermediate_op::{
 };
 use crate::NUM_CPUS;
 
+fn num_parallel_exprs(projection: &[ExprRef]) -> usize {
+    max(
+        projection
+            .iter()
+            .map(|expr| expr.has_compute())
+            .filter(|x| *x)
+            .count(),
+        1,
+    )
+}
+
 pub struct ProjectOperator {
     projection: Arc<Vec<ExprRef>>,
+    max_concurrency: usize,
+    parallel_exprs: usize,
 }
 
 impl ProjectOperator {
-    pub fn new(projection: Vec<ExprRef>) -> Self {
-        Self {
+    pub fn new(projection: Vec<ExprRef>) -> DaftResult<Self> {
+        let (max_concurrency, parallel_exprs) = Self::get_optimal_allocation(&projection)?;
+        Ok(Self {
             projection: Arc::new(projection),
-        }
+            max_concurrency,
+            parallel_exprs,
+        })
+    }
+
+    // This function is used to determine the optimal allocation of concurrency and expression parallelism
+    fn get_optimal_allocation(projection: &[ExprRef]) -> DaftResult<(usize, usize)> {
+        let resource_request = get_resource_request(projection);
+        // The number of CPUs available for the operator.
+        let available_cpus = match resource_request {
+            // If the resource request specifies a number of CPUs, the available cpus is the number of actual CPUs
+            // divided by the requested number of CPUs, clamped to (1, NUM_CPUS).
+            // E.g. if the resource request specifies 2 CPUs and NUM_CPUS is 4, the number of available cpus is 2.
+            Some(resource_request) if resource_request.num_cpus().is_some() => {
+                let requested_num_cpus = resource_request.num_cpus().unwrap();
+                if requested_num_cpus > *NUM_CPUS as f64 {
+                    Err(DaftError::ValueError(format!(
+                        "Requested {} CPUs but found only {} available",
+                        requested_num_cpus, *NUM_CPUS
+                    )))
+                } else {
+                    Ok(
+                        (*NUM_CPUS as f64 / requested_num_cpus).clamp(1.0, *NUM_CPUS as f64)
+                            as usize,
+                    )
+                }
+            }
+            _ => Ok(*NUM_CPUS),
+        }?;
+
+        let max_parallel_exprs = num_parallel_exprs(projection);
+
+        // Calculate optimal concurrency using ceiling division
+        // Example: For 128 CPUs and 60 parallel expressions:
+        // max_concurrency = (128 + 60 - 1) / 60 = 3 concurrent tasks
+        // This ensures we never exceed max_parallel_exprs per task
+        let max_concurrency = (available_cpus + max_parallel_exprs - 1) / max_parallel_exprs;
+
+        // Calculate actual parallel expressions per task using floor division
+        // Example: For 128 CPUs and 3 concurrent tasks:
+        // num_parallel_exprs = 128 / 3 = 42 parallel expressions per task
+        // This ensures even distribution across concurrent tasks
+        let num_parallel_exprs = available_cpus / max_concurrency;
+
+        Ok((max_concurrency, num_parallel_exprs))
     }
 }
 
@@ -33,9 +91,16 @@ impl IntermediateOperator for ProjectOperator {
         runtime: &RuntimeRef,
     ) -> IntermediateOpExecuteResult {
         let projection = self.projection.clone();
+        let num_parallel_exprs = self.parallel_exprs;
         runtime
             .spawn(async move {
-                let out = input.eval_expression_list(&projection)?;
+                let out = if num_parallel_exprs > 1 {
+                    input
+                        .par_eval_expression_list(&projection, num_parallel_exprs)
+                        .await?
+                } else {
+                    input.eval_expression_list(&projection)?
+                };
                 Ok((
                     state,
                     IntermediateOperatorResult::NeedMoreInput(Some(Arc::new(out))),
@@ -49,26 +114,6 @@ impl IntermediateOperator for ProjectOperator {
     }
 
     fn max_concurrency(&self) -> DaftResult<usize> {
-        let resource_request = get_resource_request(&self.projection);
-        match resource_request {
-            // If the resource request specifies a number of CPUs, the max concurrency is the number of CPUs
-            // divided by the requested number of CPUs, clamped to (1, NUM_CPUS).
-            // E.g. if the resource request specifies 2 CPUs and NUM_CPUS is 4, the max concurrency is 2.
-            Some(resource_request) if resource_request.num_cpus().is_some() => {
-                let requested_num_cpus = resource_request.num_cpus().unwrap();
-                if requested_num_cpus > *NUM_CPUS as f64 {
-                    Err(DaftError::ValueError(format!(
-                        "Requested {} CPUs but found only {} available",
-                        requested_num_cpus, *NUM_CPUS
-                    )))
-                } else {
-                    Ok(
-                        (*NUM_CPUS as f64 / requested_num_cpus).clamp(1.0, *NUM_CPUS as f64)
-                            as usize,
-                    )
-                }
-            }
-            _ => Ok(*NUM_CPUS),
-        }
+        Ok(self.max_concurrency)
     }
 }