Add crossbeam dependencies and initialize global Rayon pool

Author: kosiew

Cargo.lock

@@ -45,7 +45,10 @@ use crate::udaf::PyAggregateUDF;
 use crate::udf::PyScalarUDF;
 use crate::udtf::PyTableFunction;
 use crate::udwf::PyWindowUDF;
-use crate::utils::{get_global_ctx, spawn_and_wait, validate_pycapsule, wait_for_future};
+use crate::utils::{
+    get_global_ctx, init_global_rayon_pool, spawn_and_wait, validate_pycapsule,
+    wait_for_future,
+};
 use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef};
 use datafusion::arrow::pyarrow::PyArrowType;
 use datafusion::arrow::record_batch::RecordBatch;
@@ -313,6 +316,7 @@ impl PySessionContext {
         } else {
             SessionConfig::default().with_information_schema(true)
         };
+        init_global_rayon_pool(config.target_partitions());
         let runtime_env_builder = if let Some(c) = runtime {
             c.builder
         } else {

src/context.rs

@@ -52,8 +52,8 @@ use crate::physical_plan::PyExecutionPlan;
 use crate::record_batch::PyRecordBatchStream;
 use crate::sql::logical::PyLogicalPlan;
 use crate::utils::{
-    get_tokio_runtime, is_ipython_env, py_obj_to_scalar_value, spawn_and_wait, validate_pycapsule,
-    wait_for_future, wait_for_stream_next,
+    get_tokio_runtime, init_global_rayon_pool, is_ipython_env, py_obj_to_scalar_value,
+    spawn_and_wait, validate_pycapsule, wait_for_future, wait_for_stream_next,
 };
 use crate::{
     errors::PyDataFusionResult,
@@ -364,6 +364,7 @@ fn record_batches_to_pyarrow(
     record_batch_class: &Bound<'_, PyAny>,
     batches: Vec<RecordBatch>,
 ) -> PyResult<Vec<PyObject>> {
+    init_global_rayon_pool(std::thread::available_parallelism().map_or(1, |n| n.get()));
     let ffi_batches: Vec<(FFI_ArrowArray, FFI_ArrowSchema)> = py
         .allow_threads(|| {
             batches

src/dataframe.rs

@@ -28,6 +28,7 @@ use futures::StreamExt;
 use pyo3::prelude::*;
 use pyo3::{exceptions::PyValueError, types::PyCapsule};
 use std::{future::Future, sync::OnceLock, time::Duration};
+use rayon::ThreadPoolBuilder;
 use tokio::{runtime::Runtime, time::sleep};
 /// Utility to get the Tokio Runtime from Python
 #[inline]
@@ -59,6 +60,16 @@ pub(crate) fn get_global_ctx() -> &'static SessionContext {
     CTX.get_or_init(SessionContext::new)
 }
 
+#[inline]
+pub(crate) fn init_global_rayon_pool(num_threads: usize) {
+    static RAYON_POOL: OnceLock<()> = OnceLock::new();
+    RAYON_POOL.get_or_init(|| {
+        let _ = ThreadPoolBuilder::new()
+            .num_threads(num_threads)
+            .build_global();
+    });
+}
+
 /// Utility to collect rust futures with GIL released and respond to
 /// Python interrupts such as ``KeyboardInterrupt``. If a signal is
 /// received while the future is running, the future is aborted and the