polars_python/
utils.rs

1use std::panic::AssertUnwindSafe;
2
3use polars::frame::DataFrame;
4use polars::series::IntoSeries;
5use polars_error::PolarsResult;
6use polars_error::signals::{KeyboardInterrupt, catch_keyboard_interrupt};
7use pyo3::exceptions::PyKeyboardInterrupt;
8use pyo3::marker::Ungil;
9use pyo3::{PyErr, PyResult, Python};
10
11use crate::dataframe::PyDataFrame;
12use crate::error::PyPolarsErr;
13use crate::series::PySeries;
14use crate::timeout::{cancel_polars_timeout, schedule_polars_timeout};
15
16// was redefined because I could not get feature flags activated?
17#[macro_export]
18macro_rules! apply_method_all_arrow_series2 {
19    ($self:expr, $method:ident, $($args:expr),*) => {
20        match $self.dtype() {
21            DataType::Boolean => $self.bool().unwrap().$method($($args),*),
22            DataType::String => $self.str().unwrap().$method($($args),*),
23            DataType::UInt8 => $self.u8().unwrap().$method($($args),*),
24            DataType::UInt16 => $self.u16().unwrap().$method($($args),*),
25            DataType::UInt32 => $self.u32().unwrap().$method($($args),*),
26            DataType::UInt64 => $self.u64().unwrap().$method($($args),*),
27            DataType::UInt128 => $self.u128().unwrap().$method($($args),*),
28            DataType::Int8 => $self.i8().unwrap().$method($($args),*),
29            DataType::Int16 => $self.i16().unwrap().$method($($args),*),
30            DataType::Int32 => $self.i32().unwrap().$method($($args),*),
31            DataType::Int64 => $self.i64().unwrap().$method($($args),*),
32            DataType::Int128 => $self.i128().unwrap().$method($($args),*),
33            DataType::Float32 => $self.f32().unwrap().$method($($args),*),
34            DataType::Float64 => $self.f64().unwrap().$method($($args),*),
35            DataType::Date => $self.date().unwrap().physical().$method($($args),*),
36            DataType::Datetime(_, _) => $self.datetime().unwrap().physical().$method($($args),*),
37            DataType::List(_) => $self.list().unwrap().$method($($args),*),
38            DataType::Struct(_) => $self.struct_().unwrap().$method($($args),*),
39            dt => panic!("dtype {:?} not supported", dt)
40        }
41    }
42}
43
44/// Boilerplate for `|e| PyPolarsErr::from(e).into()`
45#[allow(unused)]
46pub(crate) fn to_py_err<E: Into<PyPolarsErr>>(e: E) -> PyErr {
47    e.into().into()
48}
49
50pub trait EnterPolarsExt {
51    /// Whenever you have a block of code in the public Python API that
52    /// (potentially) takes a long time, wrap it in enter_polars. This will
53    /// ensure we release the GIL and catch KeyboardInterrupts.
54    ///
55    /// This not only can increase performance and usability, it can avoid
56    /// deadlocks on the GIL for Python UDFs.
57    fn enter_polars<T, E, F>(self, f: F) -> PyResult<T>
58    where
59        F: Ungil + Send + FnOnce() -> Result<T, E>,
60        T: Ungil + Send,
61        E: Ungil + Send + Into<PyPolarsErr>;
62
63    /// Same as enter_polars, but wraps the result in PyResult::Ok, useful
64    /// shorthand for infallible functions.
65    #[inline(always)]
66    fn enter_polars_ok<T, F>(self, f: F) -> PyResult<T>
67    where
68        Self: Sized,
69        F: Ungil + Send + FnOnce() -> T,
70        T: Ungil + Send,
71    {
72        self.enter_polars(move || PyResult::Ok(f()))
73    }
74
75    /// Same as enter_polars, but expects a PolarsResult<DataFrame> as return
76    /// which is converted to a PyDataFrame.
77    #[inline(always)]
78    fn enter_polars_df<F>(self, f: F) -> PyResult<PyDataFrame>
79    where
80        Self: Sized,
81        F: Ungil + Send + FnOnce() -> PolarsResult<DataFrame>,
82    {
83        self.enter_polars(f).map(PyDataFrame::new)
84    }
85
86    /// Same as enter_polars, but expects a PolarsResult<S> as return which
87    /// is converted to a PySeries through S: IntoSeries.
88    #[inline(always)]
89    fn enter_polars_series<T, F>(self, f: F) -> PyResult<PySeries>
90    where
91        Self: Sized,
92        T: Ungil + Send + IntoSeries,
93        F: Ungil + Send + FnOnce() -> PolarsResult<T>,
94    {
95        self.enter_polars(f).map(|s| PySeries::new(s.into_series()))
96    }
97}
98
99impl EnterPolarsExt for Python<'_> {
100    fn enter_polars<T, E, F>(self, f: F) -> PyResult<T>
101    where
102        F: Ungil + Send + FnOnce() -> Result<T, E>,
103        T: Ungil + Send,
104        E: Ungil + Send + Into<PyPolarsErr>,
105    {
106        let timeout = schedule_polars_timeout();
107        let ret = self.detach(|| catch_keyboard_interrupt(AssertUnwindSafe(f)));
108        cancel_polars_timeout(timeout);
109        match ret {
110            Ok(Ok(ret)) => Ok(ret),
111            Ok(Err(err)) => Err(PyErr::from(err.into())),
112            Err(KeyboardInterrupt) => Err(PyKeyboardInterrupt::new_err("")),
113        }
114    }
115}
polars_python/utils.rs

polars_python/
utils.rs
