fix(python): Release the GIL in Python APIs, part 2 of 2 (#19762)

Co-authored-by: Itamar Turner-Trauring <[email protected]>

crates/polars-python/src/cloud.rs

@@ -49,9 +49,7 @@ pub fn _execute_ir_plan_with_gpu(ir_plan_ser: Vec<u8>, py: Python) -> PyResult<P
 
     // Execute the plan.
     let mut state = ExecutionState::new();
-    let df = physical_plan
-        .execute(&mut state)
-        .map_err(PyPolarsErr::from)?;
+    let df = py.allow_threads(|| physical_plan.execute(&mut state).map_err(PyPolarsErr::from))?;
 
     Ok(df.into())
 }

crates/polars-python/src/dataframe/construction.rs

@@ -52,8 +52,8 @@ impl PyDataFrame {
     }
 
     #[staticmethod]
-    pub fn from_arrow_record_batches(rb: Vec<Bound<PyAny>>) -> PyResult<Self> {
-        let df = interop::arrow::to_rust::to_rust_df(&rb)?;
+    pub fn from_arrow_record_batches(py: Python, rb: Vec<Bound<PyAny>>) -> PyResult<Self> {
+        let df = interop::arrow::to_rust::to_rust_df(py, &rb)?;
         Ok(Self::from(df))
     }
 }

crates/polars-python/src/dataframe/export.rs

@@ -79,19 +79,17 @@ impl PyDataFrame {
     }
 
     #[allow(clippy::wrong_self_convention)]
-    pub fn to_arrow(&mut self, compat_level: PyCompatLevel) -> PyResult<Vec<PyObject>> {
-        self.df.align_chunks_par();
-        Python::with_gil(|py| {
-            let pyarrow = py.import_bound("pyarrow")?;
-            let names = self.df.get_column_names_str();
+    pub fn to_arrow(&mut self, py: Python, compat_level: PyCompatLevel) -> PyResult<Vec<PyObject>> {
+        py.allow_threads(|| self.df.align_chunks_par());
+        let pyarrow = py.import_bound("pyarrow")?;
+        let names = self.df.get_column_names_str();
 
-            let rbs = self
-                .df
-                .iter_chunks(compat_level.0, true)
-                .map(|rb| interop::arrow::to_py::to_py_rb(&rb, &names, py, &pyarrow))
-                .collect::<PyResult<_>>()?;
-            Ok(rbs)
-        })
+        let rbs = self
+            .df
+            .iter_chunks(compat_level.0, true)
+            .map(|rb| interop::arrow::to_py::to_py_rb(&rb, &names, py, &pyarrow))
+            .collect::<PyResult<_>>()?;
+        Ok(rbs)
     }
 
     /// Create a `Vec` of PyArrow RecordBatch instances.
@@ -100,8 +98,8 @@ impl PyDataFrame {
     /// since those can't be converted correctly via PyArrow. The calling Python
     /// code should make sure these are not included.
     #[allow(clippy::wrong_self_convention)]
-    pub fn to_pandas(&mut self) -> PyResult<Vec<PyObject>> {
-        self.df.as_single_chunk_par();
+    pub fn to_pandas(&mut self, py: Python) -> PyResult<Vec<PyObject>> {
+        py.allow_threads(|| self.df.as_single_chunk_par());
         Python::with_gil(|py| {
             let pyarrow = py.import_bound("pyarrow")?;
             let names = self.df.get_column_names_str();
@@ -154,7 +152,7 @@ impl PyDataFrame {
         py: Python<'py>,
         requested_schema: Option<PyObject>,
     ) -> PyResult<Bound<'py, PyCapsule>> {
-        self.df.align_chunks_par();
+        py.allow_threads(|| self.df.align_chunks_par());
         dataframe_to_stream(&self.df, py)
     }
 }