polars_utils/
python_thread_pool.rs

1use pyo3::call::PyCallArgs;
2use pyo3::sync::PyOnceLock;
3use pyo3::types::{PyAnyMethods as _, PyDict};
4use pyo3::{Bound, IntoPyObject, Py, PyAny, PyResult, Python};
5
6#[derive(IntoPyObject, Debug)]
7pub struct PyThreadPool(
8    /// polars._utils.threading.PyThreadPool
9    Py<PyAny>,
10);
11
12impl<'py> IntoPyObject<'py> for &'py PyThreadPool {
13    type Output = <&'py Py<PyAny> as IntoPyObject<'py>>::Output;
14    type Target = <&'py Py<PyAny> as IntoPyObject<'py>>::Target;
15    type Error = <&'py Py<PyAny> as IntoPyObject<'py>>::Error;
16
17    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
18        IntoPyObject::into_pyobject(&self.0, py)
19    }
20}
21
22impl PyThreadPool {
23    pub fn new(num_threads: usize) -> Self {
24        return Python::attach(|py| {
25            Self(
26                py_scan_resolve_threadpool_cls(py)
27                    .bind(py)
28                    .call1((num_threads,))
29                    .map(|x| x.unbind())
30                    .unwrap(),
31            )
32        });
33
34        fn py_scan_resolve_threadpool_cls(py: Python<'_>) -> &'static Py<PyAny> {
35            static CLS: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
36
37            CLS.get_or_init(py, || {
38                py.import("polars._utils.threading")
39                    .unwrap()
40                    .getattr("PyThreadPool")
41                    .unwrap()
42                    .unbind()
43            })
44        }
45    }
46
47    pub fn new_unbounded() -> Self {
48        Self::new(usize::MAX)
49    }
50
51    pub fn new_scan_resolve_thread_pool() -> Self {
52        use std::num::NonZeroUsize;
53
54        let num_threads: usize =
55            std::env::var("POLARS_PYTHON_SCAN_RESOLVE_THREADS").map_or(128, |x| {
56                x.parse::<NonZeroUsize>()
57                    .unwrap_or_else(|_| {
58                        panic!("invalid value for POLARS_PYTHON_SCAN_RESOLVE_THREADS: {x}")
59                    })
60                    .get()
61            });
62
63        if polars_config::config().verbose() {
64            eprintln!("python scan_resolve_threadpool threads: {num_threads}")
65        }
66
67        Self::new(num_threads)
68    }
69
70    pub fn spawn_call<'a>(
71        &self,
72        py: Python<'a>,
73        function: &Py<PyAny>,
74        args: impl PyCallArgs<'a>,
75        kwargs: Option<&Bound<'a, PyDict>>,
76    ) -> PyResult<Py<PyAny>> {
77        static CLS: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
78
79        CLS.get_or_init(py, || {
80            py.import("polars._utils.threading")
81                .unwrap()
82                .getattr("FnPoolWrap")
83                .unwrap()
84                .unbind()
85        })
86        .call1(py, (function, self))?
87        .call(py, args, kwargs)
88    }
89}
polars_utils/python_thread_pool.rs

polars_utils/
python_thread_pool.rs
