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)]
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() -> Self {
24        use std::num::NonZeroUsize;
25
26        Python::attach(|py| {
27            let num_threads =
28                std::env::var("POLARS_PYTHON_SCAN_RESOLVE_THREADS").map_or(128, |x| {
29                    x.parse::<NonZeroUsize>()
30                        .unwrap_or_else(|_| {
31                            panic!("invalid value for POLARS_PYTHON_SCAN_RESOLVE_THREADS: {x}")
32                        })
33                        .get()
34                });
35
36            if polars_config::config().verbose() {
37                eprintln!("python scan_resolve_threadpool threads: {num_threads}")
38            }
39
40            return Self(
41                py_scan_resolve_threadpool_cls(py)
42                    .bind(py)
43                    .call1((num_threads,))
44                    .map(|x| x.unbind())
45                    .unwrap(),
46            );
47
48            fn py_scan_resolve_threadpool_cls(py: Python<'_>) -> &'static Py<PyAny> {
49                static CLS: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
50
51                CLS.get_or_init(py, || {
52                    py.import("polars._utils.threading")
53                        .unwrap()
54                        .getattr("PyThreadPool")
55                        .unwrap()
56                        .unbind()
57                })
58            }
59        })
60    }
61
62    pub fn spawn_call<'a>(
63        &self,
64        py: Python<'a>,
65        function: &Py<PyAny>,
66        args: impl PyCallArgs<'a>,
67        kwargs: Option<&Bound<'a, PyDict>>,
68    ) -> PyResult<Py<PyAny>> {
69        static CLS: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
70
71        CLS.get_or_init(py, || {
72            py.import("polars._utils.threading")
73                .unwrap()
74                .getattr("FnPoolWrap")
75                .unwrap()
76                .unbind()
77        })
78        .call1(py, (function, self))?
79        .call(py, args, kwargs)
80    }
81}
82
83impl Default for PyThreadPool {
84    fn default() -> Self {
85        Self::new()
86    }
87}
polars_utils/python_thread_pool.rs

polars_utils/
python_thread_pool.rs
