Python: allow_threads in collect_all

pola-rs · Oct 11, 2021 · 4d5facf · 4d5facf
1 parent 6c876bc
commit 4d5facf
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 14 deletions.
diff --git a/py-polars/src/lazy/dataframe.rs b/py-polars/src/lazy/dataframe.rs
@@ -185,13 +185,13 @@ impl PyLazyFrame {
         ldf.cache().into()
     }
 
-    pub fn collect(&self) -> PyResult<PyDataFrame> {
-        let ldf = self.ldf.clone();
-        let gil = Python::acquire_gil();
-        let py = gil.python();
+    pub fn collect(&self, py: Python) -> PyResult<PyDataFrame> {
         // if we don't allow threads and we have udfs trying to acquire the gil from different
         // threads we deadlock.
-        let df = py.allow_threads(|| ldf.collect().map_err(PyPolarsEr::from))?;
+        let df = py.allow_threads(|| {
+            let ldf = self.ldf.clone();
+            ldf.collect().map_err(PyPolarsEr::from)
+        })?;
         Ok(df.into())
     }
 

diff --git a/py-polars/src/lib.rs b/py-polars/src/lib.rs
@@ -178,16 +178,19 @@ fn ipc_schema(py: Python, py_f: PyObject) -> PyResult<PyObject> {
 }
 
 #[pyfunction]
-fn collect_all(lfs: Vec<PyLazyFrame>) -> PyResult<Vec<PyDataFrame>> {
+fn collect_all(lfs: Vec<PyLazyFrame>, py: Python) -> PyResult<Vec<PyDataFrame>> {
     use polars_core::utils::rayon::prelude::*;
-    let out = polars_core::POOL.install(|| {
-        lfs.par_iter()
-            .map(|lf| {
-                let df = lf.ldf.clone().collect()?;
-                Ok(PyDataFrame::new(df))
-            })
-            .collect::<polars_core::error::Result<Vec<_>>>()
-            .map_err(PyPolarsEr::from)
+
+    let out = py.allow_threads(|| {
+        polars_core::POOL.install(|| {
+            lfs.par_iter()
+                .map(|lf| {
+                    let df = lf.ldf.clone().collect()?;
+                    Ok(PyDataFrame::new(df))
+                })
+                .collect::<polars_core::error::Result<Vec<_>>>()
+                .map_err(PyPolarsEr::from)
+        })
     });
 
     Ok(out?)