From b826806a9bfa939849392be2994673fa83261436 Mon Sep 17 00:00:00 2001 From: Nico Mandery Date: Thu, 21 Nov 2024 18:50:40 +0100 Subject: [PATCH] release the GIL more often to allow other threads to run --- h3ronpy/CHANGES.rst | 3 +- h3ronpy/src/array.rs | 28 +++++- h3ronpy/src/op/compact.rs | 34 ++++--- h3ronpy/src/op/localij.rs | 20 ++-- h3ronpy/src/op/resolution.rs | 16 ++- h3ronpy/src/op/string.rs | 81 +++++++++------- h3ronpy/src/op/valid.rs | 4 +- h3ronpy/src/raster.rs | 13 ++- h3ronpy/src/vector.rs | 182 +++++++++++++++++++---------------- 9 files changed, 228 insertions(+), 153 deletions(-) diff --git a/h3ronpy/CHANGES.rst b/h3ronpy/CHANGES.rst index 0c97256..4068e15 100644 --- a/h3ronpy/CHANGES.rst +++ b/h3ronpy/CHANGES.rst @@ -12,7 +12,8 @@ Versioning `__. 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 index 47d98b5..4aa5b8f 100644 --- 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 for PyDirectedEdgeArray { + fn as_ref(&self) -> &DirectedEdgeIndexArray { + &self.0 + } +} + +impl<'py> FromPyObject<'py> for PyDirectedEdgeArray { + fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + 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 for PyVertexArray { + fn as_ref(&self) -> &VertexIndexArray { + &self.0 + } +} + +impl<'py> FromPyObject<'py> for PyVertexArray { + fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + Ok(Self(pyarray_to_vertexindexarray(ob)?)) + } +} diff --git a/h3ronpy/src/op/compact.rs b/h3ronpy/src/op/compact.rs index fa163b9..a77acc9 100644 --- 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 { +pub(crate) fn compact( + py: Python<'_>, + cellarray: PyCellArray, + mixed_resolutions: bool, +) -> PyResult { 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 { +pub(crate) fn uncompact( + py: Python<'_>, + cellarray: PyCellArray, + target_resolution: u8, +) -> PyResult { 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 index edb52f7..03689e0 100644 --- 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, i_array: &Bound, j_array: &Bound, @@ -59,17 +60,18 @@ pub(crate) fn localij_to_cells( let j_array = pyarray_to_native::(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, len: usize) -> PyResult { diff --git a/h3ronpy/src/op/resolution.rs b/h3ronpy/src/op/resolution.rs index b56bd4c..25b25fc 100644 --- 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 { +pub(crate) fn change_resolution( + py: Python<'_>, + cellarray: PyCellArray, + h3_resolution: u8, +) -> PyResult { 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 index 5af79f9..620bac7 100644 --- 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 { let (boxed_array, _field) = stringarray.into_inner(); - let cells = if let Some(stringarray) = boxed_array.as_any().downcast_ref::() { - CellIndexArray::parse_genericstringarray(stringarray, set_failing_to_invalid) - .into_pyresult()? - } else if let Some(stringarray) = boxed_array.as_any().downcast_ref::() { - 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::() { + CellIndexArray::parse_genericstringarray(stringarray, set_failing_to_invalid) + .into_pyresult() + } else if let Some(stringarray) = boxed_array.as_any().downcast_ref::() { + 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 { let (boxed_array, _field) = stringarray.into_inner(); - let vertexes = if let Some(utf8array) = boxed_array.as_any().downcast_ref::() { - VertexIndexArray::parse_genericstringarray(utf8array, set_failing_to_invalid) - .into_pyresult()? - } else if let Some(utf8array) = boxed_array.as_any().downcast_ref::() { - 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::() { + VertexIndexArray::parse_genericstringarray(utf8array, set_failing_to_invalid) + .into_pyresult() + } else if let Some(utf8array) = boxed_array.as_any().downcast_ref::() { + 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 { let (boxed_array, _field) = stringarray.into_inner(); - let edges = if let Some(stringarray) = boxed_array.as_any().downcast_ref::() { - DirectedEdgeIndexArray::parse_genericstringarray(stringarray, set_failing_to_invalid) - .into_pyresult()? - } else if let Some(stringarray) = boxed_array.as_any().downcast_ref::() { - 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::() { + DirectedEdgeIndexArray::parse_genericstringarray(stringarray, set_failing_to_invalid) + .into_pyresult() + } else if let Some(stringarray) = boxed_array.as_any().downcast_ref::() { + 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 index 78710b3..8126958 100644 --- a/h3ronpy/src/op/valid.rs +++ b/h3ronpy/src/op/valid.rs @@ -12,10 +12,10 @@ use crate::arrow_interop::*; fn h3index_valid(py: Python, arr: &Bound, booleanarray: bool) -> PyResult where - IX: H3IndexArrayValue, + IX: H3IndexArrayValue + Send, { let u64array = pyarray_to_uint64array(arr)?; - let validated = H3Array::::from_iter_with_validity(u64array.iter()); + let validated = py.allow_threads(|| H3Array::::from_iter_with_validity(u64array.iter())); if booleanarray { let nullbuffer = validated diff --git a/h3ronpy/src/raster.rs b/h3ronpy/src/raster.rs index 4ceea90..d4eca1e 100644 --- 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::>()), + 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)?; diff --git a/h3ronpy/src/vector.rs b/h3ronpy/src/vector.rs index 8cb4333..9696752 100644 --- a/h3ronpy/src/vector.rs +++ b/h3ronpy/src/vector.rs @@ -24,7 +24,7 @@ use pyo3::types::PyTuple; use pyo3_arrow::error::PyArrowResult; use pyo3_arrow::{PyArray, PyRecordBatch}; -use crate::array::PyCellArray; +use crate::array::{PyCellArray, PyDirectedEdgeArray, PyVertexArray}; use crate::arrow_interop::*; use crate::error::IntoPyResult; @@ -87,14 +87,12 @@ impl PyContainmentMode { #[pyfunction] #[pyo3(signature = (cellarray,))] -pub(crate) fn cells_bounds(cellarray: PyCellArray) -> PyResult> { - if let Some(rect) = cellarray.as_ref().bounding_rect() { - Python::with_gil(|py| { - Ok(Some( - PyTuple::new_bound(py, [rect.min().x, rect.min().y, rect.max().x, rect.max().y]) - .to_object(py), - )) - }) +pub(crate) fn cells_bounds(py: Python<'_>, cellarray: PyCellArray) -> PyResult> { + if let Some(rect) = py.allow_threads(|| cellarray.as_ref().bounding_rect()) { + Ok(Some( + PyTuple::new_bound(py, [rect.min().x, rect.min().y, rect.max().x, rect.max().y]) + .to_object(py), + )) } else { Ok(None) } @@ -176,6 +174,7 @@ pub(crate) fn cells_to_coordinates( #[pyfunction] #[pyo3(signature = (latarray, lngarray, resolution, radians = false))] pub(crate) fn coordinates_to_cells( + py: Python<'_>, latarray: &Bound, lngarray: &Bound, resolution: &Bound, @@ -192,22 +191,24 @@ pub(crate) fn coordinates_to_cells( let cells = if let Ok(resolution) = resolution.extract::() { let resolution = Resolution::try_from(resolution).into_pyresult()?; - latarray - .iter() - .zip(lngarray.iter()) - .map(|(lat, lng)| { - if let (Some(lat), Some(lng)) = (lat, lng) { - if radians { - LatLng::from_radians(lat, lng).into_pyresult() + py.allow_threads(|| { + latarray + .iter() + .zip(lngarray.iter()) + .map(|(lat, lng)| { + if let (Some(lat), Some(lng)) = (lat, lng) { + if radians { + LatLng::from_radians(lat, lng).into_pyresult() + } else { + LatLng::new(lat, lng).into_pyresult() + } + .map(|ll| Some(ll.to_cell(resolution))) } else { - LatLng::new(lat, lng).into_pyresult() + Ok(None) } - .map(|ll| Some(ll.to_cell(resolution))) - } else { - Ok(None) - } - }) - .collect::>()? + }) + .collect::>() + })? } else { let resarray = ResolutionArray::try_from(pyarray_to_native::(resolution)?) .into_pyresult()?; @@ -218,23 +219,25 @@ pub(crate) fn coordinates_to_cells( )); } - multizip((latarray.iter(), lngarray.iter(), resarray.iter())) - .map(|(lat, lng, res)| { - if let (Some(lat), Some(lng), Some(res)) = (lat, lng, res) { - if radians { - LatLng::from_radians(lat, lng).into_pyresult() + py.allow_threads(|| { + multizip((latarray.iter(), lngarray.iter(), resarray.iter())) + .map(|(lat, lng, res)| { + if let (Some(lat), Some(lng), Some(res)) = (lat, lng, res) { + if radians { + LatLng::from_radians(lat, lng).into_pyresult() + } else { + LatLng::new(lat, lng).into_pyresult() + } + .map(|ll| Some(ll.to_cell(res))) } else { - LatLng::new(lat, lng).into_pyresult() + Ok(None) } - .map(|ll| Some(ll.to_cell(res))) - } else { - Ok(None) - } - }) - .collect::>()? + }) + .collect::>() + })? }; - Python::with_gil(|py| h3array_to_pyarray(cells, py)) + h3array_to_pyarray(cells, py) } #[pyfunction] @@ -248,31 +251,33 @@ pub(crate) fn cells_to_wkb_polygons( let cellindexarray = cellarray.into_inner(); let use_degrees = !radians; - let out: WKBArray = if link_cells { - let mut cells = cellindexarray.iter().flatten().collect::>(); - cells.sort_unstable(); - cells.dedup(); + let out: WKBArray = py.allow_threads(|| { + if link_cells { + let mut cells = cellindexarray.iter().flatten().collect::>(); + cells.sort_unstable(); + cells.dedup(); - let geoms = dissolve(cells) - .into_pyresult()? - .into_iter() - .map(|mut poly| { - if radians { - poly.to_radians_in_place(); - } - Some(geo_types::Geometry::from(poly)) - }) - .collect::>(); - let mut builder = WKBBuilder::with_capacity(WKBCapacity::from_geometries( - geoms.iter().map(|v| v.as_ref()), - )); - builder.extend_from_iter(geoms.iter().map(|v| v.as_ref())); - builder.finish() - } else { - cellindexarray - .to_wkb_polygons(use_degrees) - .expect("wkbarray") - }; + let geoms = dissolve(cells) + .into_pyresult()? + .into_iter() + .map(|mut poly| { + if radians { + poly.to_radians_in_place(); + } + Some(geo_types::Geometry::from(poly)) + }) + .collect::>(); + let mut builder = WKBBuilder::with_capacity(WKBCapacity::from_geometries( + geoms.iter().map(|v| v.as_ref()), + )); + builder.extend_from_iter(geoms.iter().map(|v| v.as_ref())); + Ok::<_, PyErr>(builder.finish()) + } else { + Ok(cellindexarray + .to_wkb_polygons(use_degrees) + .expect("wkbarray")) + } + })?; let field = out.extension_field(); PyArray::new(out.into_array_ref(), field).to_arro3(py) @@ -285,10 +290,12 @@ pub(crate) fn cells_to_wkb_points( cellarray: PyCellArray, radians: bool, ) -> PyResult { - let out = cellarray - .as_ref() - .to_wkb_points::(!radians) - .expect("wkbarray"); + let out = py.allow_threads(|| { + cellarray + .as_ref() + .to_wkb_points::(!radians) + .expect("wkbarray") + }); let field = out.extension_field(); PyArray::new(out.into_array_ref(), field).to_arro3(py) @@ -298,12 +305,15 @@ pub(crate) fn cells_to_wkb_points( #[pyo3(signature = (vertexarray, radians = false))] pub(crate) fn vertexes_to_wkb_points( py: Python, - vertexarray: &Bound, + vertexarray: PyVertexArray, radians: bool, ) -> PyResult { - let out = pyarray_to_vertexindexarray(vertexarray)? - .to_wkb_points::(!radians) - .expect("wkbarray"); + let out = py.allow_threads(|| { + vertexarray + .as_ref() + .to_wkb_points::(!radians) + .expect("wkbarray") + }); let field = out.extension_field(); PyArray::new(out.into_array_ref(), field).to_arro3(py) @@ -313,12 +323,15 @@ pub(crate) fn vertexes_to_wkb_points( #[pyo3(signature = (array, radians = false))] pub(crate) fn directededges_to_wkb_linestrings( py: Python, - array: &Bound, + array: PyDirectedEdgeArray, radians: bool, ) -> PyResult { - let out = pyarray_to_directededgeindexarray(array)? - .to_wkb_linestrings::(!radians) - .expect("wkbarray"); + let out = py.allow_threads(|| { + array + .as_ref() + .to_wkb_linestrings::(!radians) + .expect("wkbarray") + }); let field = out.extension_field(); PyArray::new(out.into_array_ref(), field).to_arro3(py) @@ -376,12 +389,16 @@ fn generic_wkb_to_cells( let wkbarray = WKBArray::new(binarray, Default::default()); if flatten { - let cells = wkbarray.to_cellindexarray(options).into_pyresult()?; + let cells = py + .allow_threads(|| wkbarray.to_cellindexarray(options)) + .into_pyresult()?; - Python::with_gil(|py| h3array_to_pyarray(cells, py)) + h3array_to_pyarray(cells, py) } else { - let listarray: GenericListArray = - wkbarray.to_celllistarray(options).into_pyresult()?.into(); + let listarray: GenericListArray = py + .allow_threads(|| wkbarray.to_celllistarray(options)) + .into_pyresult()? + .into(); PyArray::from_array_ref(Arc::new(listarray)).to_arro3(py) } } @@ -389,19 +406,22 @@ fn generic_wkb_to_cells( #[pyfunction] #[pyo3(signature = (obj, resolution, containment_mode = None, compact = false))] pub(crate) fn geometry_to_cells( + py: Python<'_>, obj: py_geo_interface::Geometry, resolution: u8, containment_mode: Option, compact: bool, ) -> PyResult { if obj.0.is_empty() { - return Python::with_gil(|py| h3array_to_pyarray(CellIndexArray::new_null(0), py)); + return h3array_to_pyarray(CellIndexArray::new_null(0), py); } let options = get_to_cells_options(resolution, containment_mode, compact)?; - let cellindexarray = CellIndexArray::from( - h3arrow::array::from_geo::geometry_to_cells(&obj.0, &options).into_pyresult()?, - ); - Python::with_gil(|py| h3array_to_pyarray(cellindexarray, py)) + let cellindexarray = py.allow_threads(|| { + Ok::<_, PyErr>(CellIndexArray::from( + h3arrow::array::from_geo::geometry_to_cells(&obj.0, &options).into_pyresult()?, + )) + })?; + h3array_to_pyarray(cellindexarray, py) } pub fn init_vector_submodule(m: &Bound) -> PyResult<()> {