release the GIL more often to allow other threads to run

nmandery · Nov 21, 2024 · c0ebdd4 · c0ebdd4
1 parent ae3d0c9
commit c0ebdd4
Show file tree

Hide file tree

Showing 9 changed files with 228 additions and 153 deletions.
diff --git a/h3ronpy/CHANGES.rst b/h3ronpy/CHANGES.rst
@@ -12,7 +12,8 @@ Versioning <https://semver.org/spec/v2.0.0.html>`__.
 Unreleased
 ----------
 
-- Upgrade to h3o 0.7
+- Upgrade to h3o 0.7.
+- Release the GIL more often to allow other threads to run.
 - The minimum supported python version is now 3.9.
 
 0.21.1 - 2024-10-04

diff --git a/h3ronpy/src/array.rs b/h3ronpy/src/array.rs
@@ -6,7 +6,9 @@ use pyo3::prelude::*;
 use pyo3::types::{PyCapsule, PyTuple};
 use pyo3_arrow::ffi::to_array_pycapsules;
 
-use crate::arrow_interop::pyarray_to_cellindexarray;
+use crate::arrow_interop::{
+    pyarray_to_cellindexarray, pyarray_to_directededgeindexarray, pyarray_to_vertexindexarray,
+};
 use crate::resolution::PyResolution;
 
 #[pyclass(name = "CellArray")]
@@ -89,6 +91,18 @@ impl PyDirectedEdgeArray {
     }
 }
 
+impl AsRef<DirectedEdgeIndexArray> for PyDirectedEdgeArray {
+    fn as_ref(&self) -> &DirectedEdgeIndexArray {
+        &self.0
+    }
+}
+
+impl<'py> FromPyObject<'py> for PyDirectedEdgeArray {
+    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
+        Ok(Self(pyarray_to_directededgeindexarray(ob)?))
+    }
+}
+
 #[pyclass(name = "VertexArray")]
 pub struct PyVertexArray(VertexIndexArray);
 
@@ -117,3 +131,15 @@ impl PyVertexArray {
         Self(self.0.slice(offset, length))
     }
 }
+
+impl AsRef<VertexIndexArray> for PyVertexArray {
+    fn as_ref(&self) -> &VertexIndexArray {
+        &self.0
+    }
+}
+
+impl<'py> FromPyObject<'py> for PyVertexArray {
+    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
+        Ok(Self(pyarray_to_vertexindexarray(ob)?))
+    }
+}
diff --git a/h3ronpy/src/op/compact.rs b/h3ronpy/src/op/compact.rs
@@ -8,22 +8,34 @@ use crate::error::IntoPyResult;
 
 #[pyfunction]
 #[pyo3(signature = (cellarray, mixed_resolutions = false))]
-pub(crate) fn compact(cellarray: PyCellArray, mixed_resolutions: bool) -> PyResult<PyObject> {
+pub(crate) fn compact(
+    py: Python<'_>,
+    cellarray: PyCellArray,
+    mixed_resolutions: bool,
+) -> PyResult<PyObject> {
     let cellindexarray = cellarray.into_inner();
-    let compacted = if mixed_resolutions {
-        cellindexarray.compact_mixed_resolutions()
-    } else {
-        cellindexarray.compact()
-    }
-    .into_pyresult()?;
+    let compacted = py
+        .allow_threads(|| {
+            if mixed_resolutions {
+                cellindexarray.compact_mixed_resolutions()
+            } else {
+                cellindexarray.compact()
+            }
+        })
+        .into_pyresult()?;
 
-    Python::with_gil(|py| h3array_to_pyarray(compacted, py))
+    h3array_to_pyarray(compacted, py)
 }
 
 #[pyfunction]
 #[pyo3(signature = (cellarray, target_resolution))]
-pub(crate) fn uncompact(cellarray: PyCellArray, target_resolution: u8) -> PyResult<PyObject> {
+pub(crate) fn uncompact(
+    py: Python<'_>,
+    cellarray: PyCellArray,
+    target_resolution: u8,
+) -> PyResult<PyObject> {
     let target_resolution = Resolution::try_from(target_resolution).into_pyresult()?;
-    let out = cellarray.into_inner().uncompact(target_resolution);
-    Python::with_gil(|py| h3array_to_pyarray(out, py))
+    let cellarray = cellarray.into_inner();
+    let out = py.allow_threads(|| cellarray.uncompact(target_resolution));
+    h3array_to_pyarray(out, py)
 }
diff --git a/h3ronpy/src/op/localij.rs b/h3ronpy/src/op/localij.rs
@@ -50,6 +50,7 @@ pub(crate) fn cells_to_localij(
 #[pyfunction]
 #[pyo3(signature = (anchor, i_array, j_array, set_failing_to_invalid = false))]
 pub(crate) fn localij_to_cells(
+    py: Python<'_>,
     anchor: &Bound<PyAny>,
     i_array: &Bound<PyAny>,
     j_array: &Bound<PyAny>,
@@ -59,17 +60,18 @@ pub(crate) fn localij_to_cells(
     let j_array = pyarray_to_native::<Int32Array>(j_array)?;
     let anchorarray = get_anchor_array(anchor, i_array.len())?;
 
-    let localij_arrays = LocalIJArrays::try_new(anchorarray, i_array, j_array).into_pyresult()?;
+    let cellarray = py.allow_threads(|| {
+        let localij_arrays =
+            LocalIJArrays::try_new(anchorarray, i_array, j_array).into_pyresult()?;
 
-    let cellarray = if set_failing_to_invalid {
-        localij_arrays
-            .to_cells_failing_to_invalid()
-            .into_pyresult()?
-    } else {
-        localij_arrays.to_cells().into_pyresult()?
-    };
+        if set_failing_to_invalid {
+            localij_arrays.to_cells_failing_to_invalid().into_pyresult()
+        } else {
+            localij_arrays.to_cells().into_pyresult()
+        }
+    })?;
 
-    Python::with_gil(|py| h3array_to_pyarray(cellarray, py))
+    h3array_to_pyarray(cellarray, py)
 }
 
 fn get_anchor_array(anchor: &Bound<PyAny>, len: usize) -> PyResult<CellIndexArray> {

diff --git a/h3ronpy/src/op/resolution.rs b/h3ronpy/src/op/resolution.rs
@@ -14,14 +14,20 @@ use crate::error::IntoPyResult;
 use crate::DEFAULT_CELL_COLUMN_NAME;
 
 #[pyfunction]
-pub(crate) fn change_resolution(cellarray: PyCellArray, h3_resolution: u8) -> PyResult<PyObject> {
+pub(crate) fn change_resolution(
+    py: Python<'_>,
+    cellarray: PyCellArray,
+    h3_resolution: u8,
+) -> PyResult<PyObject> {
     let cellindexarray = cellarray.into_inner();
     let h3_resolution = Resolution::try_from(h3_resolution).into_pyresult()?;
-    let out = cellindexarray
-        .change_resolution(h3_resolution)
-        .into_pyresult()?;
+    let out = py.allow_threads(|| {
+        cellindexarray
+            .change_resolution(h3_resolution)
+            .into_pyresult()
+    })?;
 
-    Python::with_gil(|py| h3array_to_pyarray(out, py))
+    h3array_to_pyarray(out, py)
 }
 
 #[pyfunction]

diff --git a/h3ronpy/src/op/string.rs b/h3ronpy/src/op/string.rs
@@ -14,67 +14,76 @@ use crate::error::IntoPyResult;
 #[pyfunction]
 #[pyo3(signature = (stringarray, set_failing_to_invalid = false))]
 pub(crate) fn cells_parse(
+    py: Python<'_>,
     stringarray: PyArray,
     set_failing_to_invalid: bool,
 ) -> PyResult<PyObject> {
     let (boxed_array, _field) = stringarray.into_inner();
-    let cells = if let Some(stringarray) = boxed_array.as_any().downcast_ref::<StringArray>() {
-        CellIndexArray::parse_genericstringarray(stringarray, set_failing_to_invalid)
-            .into_pyresult()?
-    } else if let Some(stringarray) = boxed_array.as_any().downcast_ref::<LargeStringArray>() {
-        CellIndexArray::parse_genericstringarray(stringarray, set_failing_to_invalid)
-            .into_pyresult()?
-    } else {
-        return Err(PyValueError::new_err(
-            "unsupported array type to parse cells from",
-        ));
-    };
+    let cells = py.allow_threads(|| {
+        if let Some(stringarray) = boxed_array.as_any().downcast_ref::<StringArray>() {
+            CellIndexArray::parse_genericstringarray(stringarray, set_failing_to_invalid)
+                .into_pyresult()
+        } else if let Some(stringarray) = boxed_array.as_any().downcast_ref::<LargeStringArray>() {
+            CellIndexArray::parse_genericstringarray(stringarray, set_failing_to_invalid)
+                .into_pyresult()
+        } else {
+            Err(PyValueError::new_err(
+                "unsupported array type to parse cells from",
+            ))
+        }
+    })?;
 
-    Python::with_gil(|py| h3array_to_pyarray(cells, py))
+    h3array_to_pyarray(cells, py)
 }
 
 #[pyfunction]
 #[pyo3(signature = (stringarray, set_failing_to_invalid = false))]
 pub(crate) fn vertexes_parse(
+    py: Python<'_>,
     stringarray: PyArray,
     set_failing_to_invalid: bool,
 ) -> PyResult<PyObject> {
     let (boxed_array, _field) = stringarray.into_inner();
-    let vertexes = if let Some(utf8array) = boxed_array.as_any().downcast_ref::<StringArray>() {
-        VertexIndexArray::parse_genericstringarray(utf8array, set_failing_to_invalid)
-            .into_pyresult()?
-    } else if let Some(utf8array) = boxed_array.as_any().downcast_ref::<LargeStringArray>() {
-        VertexIndexArray::parse_genericstringarray(utf8array, set_failing_to_invalid)
-            .into_pyresult()?
-    } else {
-        return Err(PyValueError::new_err(
-            "unsupported array type to parse vertexes from",
-        ));
-    };
+    let vertexes = py.allow_threads(|| {
+        if let Some(utf8array) = boxed_array.as_any().downcast_ref::<StringArray>() {
+            VertexIndexArray::parse_genericstringarray(utf8array, set_failing_to_invalid)
+                .into_pyresult()
+        } else if let Some(utf8array) = boxed_array.as_any().downcast_ref::<LargeStringArray>() {
+            VertexIndexArray::parse_genericstringarray(utf8array, set_failing_to_invalid)
+                .into_pyresult()
+        } else {
+            Err(PyValueError::new_err(
+                "unsupported array type to parse vertexes from",
+            ))
+        }
+    })?;
 
-    Python::with_gil(|py| h3array_to_pyarray(vertexes, py))
+    h3array_to_pyarray(vertexes, py)
 }
 
 #[pyfunction]
 #[pyo3(signature = (stringarray, set_failing_to_invalid = false))]
 pub(crate) fn directededges_parse(
+    py: Python<'_>,
     stringarray: PyArray,
     set_failing_to_invalid: bool,
 ) -> PyResult<PyObject> {
     let (boxed_array, _field) = stringarray.into_inner();
-    let edges = if let Some(stringarray) = boxed_array.as_any().downcast_ref::<StringArray>() {
-        DirectedEdgeIndexArray::parse_genericstringarray(stringarray, set_failing_to_invalid)
-            .into_pyresult()?
-    } else if let Some(stringarray) = boxed_array.as_any().downcast_ref::<LargeStringArray>() {
-        DirectedEdgeIndexArray::parse_genericstringarray(stringarray, set_failing_to_invalid)
-            .into_pyresult()?
-    } else {
-        return Err(PyValueError::new_err(
-            "unsupported array type to parse directededges from",
-        ));
-    };
+    let edges = py.allow_threads(|| {
+        if let Some(stringarray) = boxed_array.as_any().downcast_ref::<StringArray>() {
+            DirectedEdgeIndexArray::parse_genericstringarray(stringarray, set_failing_to_invalid)
+                .into_pyresult()
+        } else if let Some(stringarray) = boxed_array.as_any().downcast_ref::<LargeStringArray>() {
+            DirectedEdgeIndexArray::parse_genericstringarray(stringarray, set_failing_to_invalid)
+                .into_pyresult()
+        } else {
+            Err(PyValueError::new_err(
+                "unsupported array type to parse directededges from",
+            ))
+        }
+    })?;
 
-    Python::with_gil(|py| h3array_to_pyarray(edges, py))
+    h3array_to_pyarray(edges, py)
 }
 
 #[pyfunction]

diff --git a/h3ronpy/src/op/valid.rs b/h3ronpy/src/op/valid.rs
@@ -12,10 +12,10 @@ use crate::arrow_interop::*;
 
 fn h3index_valid<IX>(py: Python, arr: &Bound<PyAny>, booleanarray: bool) -> PyResult<PyObject>
 where
-    IX: H3IndexArrayValue,
+    IX: H3IndexArrayValue + Send,
 {
     let u64array = pyarray_to_uint64array(arr)?;
-    let validated = H3Array::<IX>::from_iter_with_validity(u64array.iter());
+    let validated = py.allow_threads(|| H3Array::<IX>::from_iter_with_validity(u64array.iter()));
 
     if booleanarray {
         let nullbuffer = validated

diff --git a/h3ronpy/src/raster.rs b/h3ronpy/src/raster.rs
@@ -153,16 +153,15 @@ macro_rules! make_raster_to_h3_variant {
             nodata_value: Option<$dtype>,
         ) -> PyResult<(PyObject, PyObject)> {
             let arr = np_array.as_array();
-            let (values, cells) = raster_to_h3(
+            let (values, cells) = py.allow_threads(|| raster_to_h3(
                 &arr,
                 transform,
                 &nodata_value,
                 h3_resolution,
                 axis_order_str,
                 compact,
-            )?;
+            ).map(|(values, cells)| (<$array_dtype>::from(values), cells)))?;
 
-            let values = <$array_dtype>::from(values);
             let values = PyArray::from_array_ref(Arc::new(values)).to_arro3(py)?;
             let cells = h3array_to_pyarray(CellIndexArray::from(cells), py)?;
 
@@ -188,17 +187,17 @@ macro_rules! make_raster_to_h3_float_variant {
             // create a copy with the values wrapped in ordered floats to
             // support the internal hashing
             let of_arr = arr.map(|v| OrderedFloat::from(*v));
-            let (values, cells) = raster_to_h3(
+            let (values, cells) = py.allow_threads(|| raster_to_h3(
                 &of_arr.view(),
                 transform,
                 &nodata_value.map(OrderedFloat::from),
                 h3_resolution,
                 axis_order_str,
                 compact,
-            )?;
+            ).map(|(values, cells)| (
+                <$array_dtype>::from(values.into_iter().map(|v| v.into_inner()).collect::<Vec<$dtype>>()),
+                cells)))?;
 
-            let values: Vec<$dtype> = values.into_iter().map(|v| v.into_inner()).collect();
-            let values = <$array_dtype>::from(values);
             let values = PyArray::from_array_ref(Arc::new(values)).to_arro3(py)?;
             let cells = h3array_to_pyarray(CellIndexArray::from(cells), py)?;