release python gil for arrow and arrow2

sfu-db · May 7, 2024 · 240716f · 240716f
1 parent bec5b1c
commit 240716f
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 6 deletions.
diff --git a/connectorx-python/src/arrow.rs b/connectorx-python/src/arrow.rs
@@ -16,9 +16,13 @@ pub fn write_arrow<'py>(
     origin_query: Option<String>,
     queries: &[CXQuery<String>],
 ) -> Bound<'py, PyAny> {
-    let destination = get_arrow(source_conn, origin_query, queries)?;
-    let rbs = destination.arrow()?;
-    let ptrs = to_ptrs(rbs);
+    let ptrs = py.allow_threads(
+        || -> Result<(Vec<String>, Vec<Vec<(uintptr_t, uintptr_t)>>), ConnectorXPythonError> {
+            let destination = get_arrow(source_conn, origin_query, queries)?;
+            let rbs = destination.arrow()?;
+            Ok(to_ptrs(rbs))
+        },
+    )?;
     let obj: PyObject = ptrs.into_py(py);
     obj.into_bound(py)
 }

diff --git a/connectorx-python/src/arrow2.rs b/connectorx-python/src/arrow2.rs
@@ -20,9 +20,13 @@ pub fn write_arrow<'py>(
     origin_query: Option<String>,
     queries: &[CXQuery<String>],
 ) -> Bound<'py, PyAny> {
-    let destination = get_arrow2(source_conn, origin_query, queries)?;
-    let (rbs, schema) = destination.arrow()?;
-    let ptrs = to_ptrs(rbs, schema);
+    let ptrs = py.allow_threads(
+        || -> Result<(Vec<String>, Vec<Vec<(uintptr_t, uintptr_t)>>), ConnectorXPythonError> {
+            let destination = get_arrow2(source_conn, origin_query, queries)?;
+            let (rbs, schema) = destination.arrow()?;
+            Ok(to_ptrs(rbs, schema))
+        },
+    )?;
     let obj: PyObject = ptrs.into_py(py);
     obj.into_bound(py)
 }