Repository: geopandas/geopandas
Pull request #3142 (2024-01-15 08:03:17)

TST: update extension array tests for upstream pandas changes


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1f88c9fa..fadd1dd8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -25,6 +25,8 @@ New methods:
 
 New features and improvements:
 
+- Added ``predicate="dwithin"`` option and ``distance`` argument to the ``sindex.query()`` method
+ and ``sjoin`` (#2882).
 - GeoSeries and GeoDataFrame `__repr__` now trims trailing zeros for a more readable
   output (#3087).
 - Add `on_invalid` parameter to `from_wkt` and `from_wkb` (#3110).
diff --git a/geopandas/_compat.py b/geopandas/_compat.py
index 4f28d1f0..f254bd64 100644
--- a/geopandas/_compat.py
+++ b/geopandas/_compat.py
@@ -23,6 +23,7 @@ PANDAS_GE_22 = Version(pd.__version__) >= Version("2.2.0.dev0")
 
 
 GEOS_GE_390 = shapely.geos.geos_version >= (3, 9, 0)
+GEOS_GE_310 = shapely.geos.geos_version >= (3, 10, 0)
 
 
 def import_optional_dependency(name: str, extra: str = ""):
diff --git a/geopandas/sindex.py b/geopandas/sindex.py
index da5f2443..34ba0650 100644
--- a/geopandas/sindex.py
+++ b/geopandas/sindex.py
@@ -5,6 +5,12 @@ import shapely
 from shapely.geometry.base import BaseGeometry
 
 from . import array, geoseries
+from . import _compat as compat
+
+PREDICATES = {p.name for p in shapely.strtree.BinaryPredicate} | {None}
+
+if compat.GEOS_GE_310:
+    PREDICATES.update(["dwithin"])
 
 
 class SpatialIndex:
@@ -44,11 +50,11 @@ class SpatialIndex:
         >>> s = geopandas.GeoSeries([Point(0, 0), Point(1, 1)])
         >>> s.sindex.valid_query_predicates  # doctest: +SKIP
         {None, "contains", "contains_properly", "covered_by", "covers", \
-"crosses", "intersects", "overlaps", "touches", "within"}
+"crosses", "dwithin", "intersects", "overlaps", "touches", "within"}
         """
-        return {p.name for p in shapely.strtree.BinaryPredicate} | {None}
+        return PREDICATES
 
-    def query(self, geometry, predicate=None, sort=False):
+    def query(self, geometry, predicate=None, sort=False, distance=None):
         """
         Return the integer indices of all combinations of each input geometry
         and tree geometries where the bounding box of each input geometry
@@ -67,6 +73,8 @@ class SpatialIndex:
         to those that meet the predicate when comparing the input geometry to
         the tree geometry: ``predicate(geometry, tree_geometry)``.
 
+        The 'dwithin' predicate requires GEOS >= 3.10.
+
         Bounding boxes are limited to two dimensions and are axis-aligned
         (equivalent to the ``bounds`` property of a geometry); any Z values
         present in input geometries are ignored when querying the tree.
@@ -83,7 +91,7 @@ class SpatialIndex:
             iterables (GeoSeries, GeometryArray) or a numpy array of Shapely
             geometries.
         predicate : {None, "contains", "contains_properly", "covered_by", "covers", \
-"crosses", "intersects", "overlaps", "touches", "within"}, optional
+"crosses", "intersects", "overlaps", "touches", "within", "dwithin"}, optional
             If predicate is provided, the input geometries are tested
             using the predicate function against each item in the tree
             whose extent intersects the envelope of the input geometry:
@@ -97,6 +105,10 @@ class SpatialIndex:
             as the secondary key.
             If False, no additional sorting is applied (results are often
             sorted but there is no guarantee).
+        distance : number or array_like, optional
+            Distances around each input geometry within which to query the tree for
+            the 'dwithin' predicate. If array_like, shape must be broadcastable to shape
+            of geometry. Required if ``predicate='dwithin'``.
 
         Returns
         -------
@@ -151,6 +163,12 @@ class SpatialIndex:
         array([[0],
                [3]])
 
+        >>> s.sindex.query(box(1, 1, 3, 3), predicate="dwithin", distance=0)
+        array([1, 2, 3])
+
+        >>> s.sindex.query(box(1, 1, 3, 3), predicate="dwithin", distance=2)
+        array([0, 1, 2, 3, 4])
+
         Notes
         -----
         In the context of a spatial join, input geometries are the "left"
@@ -161,14 +179,36 @@ class SpatialIndex:
         optional predicate are returned.
         """
         if predicate not in self.valid_query_predicates:
+            if predicate == "dwithin":
+                raise ValueError("predicate = 'dwithin' requires GEOS >= 3.10.0")
+
             raise ValueError(
-                "Got `predicate` = `{}`; ".format(predicate)
+                "Got predicate='{}'; ".format(predicate)
                 + "`predicate` must be one of {}".format(self.valid_query_predicates)
             )
 
+        # distance argument requirement of predicate `dwithin`
+        # and only valid for predicate `dwithin`
+        kwargs = {}
+        if predicate == "dwithin":
+            if distance is None:
+                # the distance parameter is needed
+                raise ValueError(
+                    "'distance' parameter is required for 'dwithin' predicate"
+                )
+            # add distance to kwargs
+            kwargs["distance"] = distance
+
+        elif distance is not None:
+            # distance parameter is invalid
+            raise ValueError(
+                "'distance' parameter is only supported in combination with "
+                "'dwithin' predicate"
+            )
+
         geometry = self._as_geometry_array(geometry)
 
-        indices = self._tree.query(geometry, predicate=predicate)
+        indices = self._tree.query(geometry, predicate=predicate, **kwargs)
 
         if sort:
             if indices.ndim == 1:
@@ -209,7 +249,7 @@ class SpatialIndex:
         else:
             return np.asarray(geometry)
 
-    def query_bulk(self, geometry, predicate=None, sort=False):
+    def query_bulk(self, geometry, predicate=None, sort=False, distance=None):
         """
         DEPRECATED: use `query` instead.
 
@@ -283,7 +323,7 @@ class SpatialIndex:
             FutureWarning,
             stacklevel=2,
         )
-        return self.query(geometry, predicate=predicate, sort=sort)
+        return self.query(geometry, predicate=predicate, sort=sort, distance=distance)
 
     def nearest(
         self,
diff --git a/geopandas/tools/sjoin.py b/geopandas/tools/sjoin.py
index 5da51223..f717f47f 100644
--- a/geopandas/tools/sjoin.py
+++ b/geopandas/tools/sjoin.py
@@ -15,6 +15,7 @@ def sjoin(
     predicate="intersects",
     lsuffix="left",
     rsuffix="right",
+    distance=None,
     **kwargs,
 ):
     """Spatial join of two GeoDataFrames.
@@ -42,6 +43,11 @@ def sjoin(
         Suffix to apply to overlapping column names (left GeoDataFrame).
     rsuffix : string, default 'right'
         Suffix to apply to overlapping column names (right GeoDataFrame).
+    distance : number or array_like, optional
+        Distance(s) around each input geometry within which to query the tree
+        for the 'dwithin' predicate. If array_like, must be
+        one-dimesional with length equal to length of left GeoDataFrame.
+        Required if ``predicate='dwithin'``.
 
     Examples
     --------
@@ -117,7 +123,7 @@ def sjoin(
 
     _basic_checks(left_df, right_df, how, lsuffix, rsuffix)
 
-    indices = _geom_predicate_query(left_df, right_df, predicate)
+    indices = _geom_predicate_query(left_df, right_df, predicate, distance)
 
     joined = _frame_join(indices, left_df, right_df, how, lsuffix, rsuffix)
 
@@ -174,7 +180,7 @@ def _basic_checks(left_df, right_df, how, lsuffix, rsuffix):
         )
 
 
-def _geom_predicate_query(left_df, right_df, predicate):
+def _geom_predicate_query(left_df, right_df, predicate, distance):
     """Compute geometric comparisons and get matching indices.
 
     Parameters
@@ -213,7 +219,9 @@ def _geom_predicate_query(left_df, right_df, predicate):
             input_geoms = left_df.geometry
 
     if sindex:
-        l_idx, r_idx = sindex.query(input_geoms, predicate=predicate, sort=False)
+        l_idx, r_idx = sindex.query(
+            input_geoms, predicate=predicate, sort=False, distance=distance
+        )
         indices = pd.DataFrame({"_key_left": l_idx, "_key_right": r_idx})
     else:
         # when sindex is empty / has no valid geometries


Test patch:
diff --git a/geopandas/tests/test_extension_array.py b/geopandas/tests/test_extension_array.py
index 7bd6ad44..7b91d64f 100644
--- a/geopandas/tests/test_extension_array.py
+++ b/geopandas/tests/test_extension_array.py
@@ -24,7 +24,7 @@ import shapely.geometry
 from shapely.geometry import Point
 
 from geopandas.array import GeometryArray, GeometryDtype, from_shapely
-from geopandas._compat import PANDAS_GE_15
+from geopandas._compat import PANDAS_GE_15, PANDAS_GE_22
 
 import pytest
 
@@ -400,8 +400,20 @@ class TestMissing(extension_tests.BaseMissingTests):
     def test_fillna_no_op_returns_copy(self, data):
         pass
 
+    @pytest.mark.skip("fillna method not supported")
+    def test_ffill_limit_area(
+        self, data_missing, limit_area, input_ilocs, expected_ilocs
+    ):
+        pass
+
+
+if PANDAS_GE_22:
+    from pandas.tests.extension.base import BaseReduceTests
+else:
+    from pandas.tests.extension.base import BaseNoReduceTests as BaseReduceTests
+
 
-class TestReduce(extension_tests.BaseNoReduceTests):
+class TestReduce(BaseReduceTests):
     @pytest.mark.skip("boolean reduce (any/all) tested in test_pandas_methods")
     def test_reduce_series_boolean(self):
         pass
diff --git a/geopandas/tests/test_geodataframe.py b/geopandas/tests/test_geodataframe.py
index a325aff9..85cb3c56 100644
--- a/geopandas/tests/test_geodataframe.py
+++ b/geopandas/tests/test_geodataframe.py
@@ -8,7 +8,7 @@ import pandas as pd
 
 from pyproj import CRS
 from pyproj.exceptions import CRSError
-from shapely.geometry import Point, Polygon
+from shapely.geometry import Point, Polygon, box
 
 import geopandas
 from geopandas import GeoDataFrame, GeoSeries, points_from_xy, read_file
@@ -17,6 +17,7 @@ from geopandas.array import GeometryArray, GeometryDtype, from_shapely
 from geopandas.testing import assert_geodataframe_equal, assert_geoseries_equal
 from geopandas.tests.util import PACKAGE_DIR, validate_boro_df
 from pandas.testing import assert_frame_equal, assert_index_equal, assert_series_equal
+import geopandas._compat as compat
 import pytest
 
 
@@ -992,6 +993,26 @@ class TestDataFrame:
         result = left.sjoin(right, how=how, predicate=predicate)
         assert_geodataframe_equal(result, expected)
 
+    @pytest.mark.parametrize("how", ["left", "inner", "right"])
+    @pytest.mark.parametrize("distance", [0, 3])
+    @pytest.mark.skipif(
+        not compat.GEOS_GE_310,
+        reason="`dwithin` requires GEOS 3.10",
+    )
+    def test_sjoin_dwithin(self, how, distance):
+        """
+        Basic test for predicate='dwithin' availability of the GeoDataFrame method.
+        Other sjoin tests are located in /tools/tests/test_sjoin.py
+        """
+        left = GeoDataFrame(geometry=points_from_xy([0, 1, 2], [0, 1, 1]))
+        right = GeoDataFrame(geometry=[box(0, 0, 1, 1)])
+
+        expected = geopandas.sjoin(
+            left, right, how=how, predicate="dwithin", distance=distance
+        )
+        result = left.sjoin(right, how=how, predicate="dwithin", distance=distance)
+        assert_geodataframe_equal(result, expected)
+
     @pytest.mark.parametrize("how", ["left", "inner", "right"])
     @pytest.mark.parametrize("max_distance", [None, 1])
     @pytest.mark.parametrize("distance_col", [None, "distance"])
diff --git a/geopandas/tests/test_pandas_methods.py b/geopandas/tests/test_pandas_methods.py
index 34bd9ddc..8408f525 100644
--- a/geopandas/tests/test_pandas_methods.py
+++ b/geopandas/tests/test_pandas_methods.py
@@ -710,7 +710,7 @@ def test_groupby_metadata(crs, geometry_name):
 
     expected = (
         df.take(take_indices)
-        .set_index("value2", drop=False, append=True)
+        .set_index("value2", drop=compat.PANDAS_GE_22, append=True)
         .swaplevel()
         .rename(columns={"value1": "value1_left"})
         .assign(value1_right=value_right)
diff --git a/geopandas/tests/test_sindex.py b/geopandas/tests/test_sindex.py
index 768ab817..51fb5dc0 100644
--- a/geopandas/tests/test_sindex.py
+++ b/geopandas/tests/test_sindex.py
@@ -380,6 +380,108 @@ class TestShapelyInterface:
         with pytest.raises(TypeError):
             self.df.sindex.query("notavalidgeom")
 
+    @pytest.mark.skipif(not compat.GEOS_GE_310, reason="Requires GEOS 3.10")
+    @pytest.mark.parametrize(
+        "distance, test_geom, expected",
+        (
+            # bounds don't intersect and not within distance=0
+            (
+                0,
+                box(9.0, 9.0, 9.9, 9.9),
+                [],
+            ),
+            # bounds don't intersect but is within distance=1
+            (
+                1,
+                box(9.0, 9.0, 9.9, 9.9),
+                [5],
+            ),
+            # within 1-D absolute distance in both axes, but not euclidean distance
+            (
+                0.5,
+                Point(0.5, 0.5),
+                [],
+            ),
+            # same as before but within euclidean distance
+            (
+                sqrt(2 * 0.5**2) + 1e-9,
+                Point(0.5, 0.5),
+                [0, 1],
+            ),
+            # less than euclidean distance between points, multi-object
+            (
+                sqrt(2) - 1e-9,
+                [
+                    Polygon([(0, 0), (1, 0), (1, 1)]),
+                    Polygon([(1, 1), (2, 1), (2, 2)]),
+                ],  # multi-object test
+                [[0, 0, 1, 1], [0, 1, 1, 2]],
+            ),
+            # more than euclidean distance between points, multi-object
+            (
+                sqrt(2) + 1e-9,
+                [
+                    Polygon([(0, 0), (1, 0), (1, 1)]),
+                    Polygon([(1, 1), (2, 1), (2, 2)]),
+                ],
+                [[0, 0, 0, 1, 1, 1, 1], [0, 1, 2, 0, 1, 2, 3]],
+            ),
+            # distance is array-like, broadcastable to geometry
+            (
+                [2, 10],
+                [Point(0.5, 0.5), Point(1, 1)],
+                [[0, 0, 1, 1, 1, 1, 1], [0, 1, 0, 1, 2, 3, 4]],
+            ),
+        ),
+    )
+    def test_query_dwithin(self, distance, test_geom, expected):
+        """Tests the `query` method with predicates that require keyword arguments."""
+        res = self.df.sindex.query(test_geom, predicate="dwithin", distance=distance)
+        assert_array_equal(res, expected)
+
+    @pytest.mark.skipif(not compat.GEOS_GE_310, reason="Requires GEOS 3.10")
+    def test_dwithin_no_distance(self):
+        """Tests the `query` method with keyword arguments that are
+        invalid for certain predicates."""
+        with pytest.raises(
+            ValueError, match="'distance' parameter is required for 'dwithin' predicate"
+        ):
+            self.df.sindex.query(Point(0, 0), predicate="dwithin")
+
+    @pytest.mark.parametrize(
+        "predicate",
+        [
+            None,
+            "contains",
+            "contains_properly",
+            "covered_by",
+            "covers",
+            "crosses",
+            "intersects",
+            "overlaps",
+            "touches",
+            "within",
+        ],
+    )
+    def test_query_distance_invalid(self, predicate):
+        """Tests the `query` method with keyword arguments that are
+        invalid for certain predicates."""
+        msg = "'distance' parameter is only supported in combination with 'dwithin'"
+        with pytest.raises(ValueError, match=msg):
+            self.df.sindex.query(Point(0, 0), predicate=predicate, distance=0)
+
+    @pytest.mark.skipif(
+        compat.GEOS_GE_310, reason="Test for 'dwithin'-incompatible versions of GEOS"
+    )
+    def test_dwithin_requirements(self):
+        """Tests whether a ValueError is raised when trying to use dwithin with
+        incompatible versions of shapely or pyGEOS
+        """
+        with pytest.raises(
+            ValueError, match="predicate = 'dwithin' requires GEOS >= 3.10.0"
+        ):
+            self.df.sindex.query(Point(0, 0), predicate="dwithin", distance=0)
+
     @pytest.mark.parametrize(
         "test_geom, expected_value",
         [
diff --git a/geopandas/tools/tests/test_sjoin.py b/geopandas/tools/tests/test_sjoin.py
index 3ddfde50..bc156c72 100644
--- a/geopandas/tools/tests/test_sjoin.py
+++ b/geopandas/tools/tests/test_sjoin.py
@@ -362,6 +362,61 @@ class TestSpatialJoin:
 
         assert_frame_equal(res, exp, check_index_type=False)
 
+    @pytest.mark.parametrize("how", ["inner"])
+    @pytest.mark.parametrize(
+        "geo_left, geo_right, expected_left, expected_right, distance",
+        [
+            (
+                # Distance is number, 2x1
+                [Point(0, 0), Point(1, 1)],
+                [Point(1, 1)],
+                [0, 1],
+                [0, 0],
+                math.sqrt(2),
+            ),
+            # Distance is number, 2x2
+            (
+                [Point(0, 0), Point(1, 1)],
+                [Point(0, 0), Point(1, 1)],
+                [0, 1, 0, 1],
+                [0, 0, 1, 1],
+                math.sqrt(2),
+            ),
+            # Distance is array, matches len(left)
+            (
+                [Point(0, 0), Point(0, 0), Point(-1, -1)],
+                [Point(1, 1)],
+                [1, 2],
+                [0, 0],
+                [0, math.sqrt(2), math.sqrt(8)],
+            ),
+            # Distance is np.array, matches len(left),
+            # inner join sorts the right GeoDataFrame
+            (
+                [Point(0, 0), Point(0, 0), Point(-1, -1)],
+                [Point(1, 1), Point(0.5, 0.5)],
+                [1, 2, 1, 2],
+                [1, 1, 0, 0],
+                np.array([0, math.sqrt(2), math.sqrt(8)]),
+            ),
+        ],
+    )
+    def test_sjoin_dwithin(
+        self,
+        geo_left,
+        geo_right,
+        expected_left: Sequence[int],
+        expected_right: Sequence[int],
+        distance,
+        how,
+    ):
+        left = geopandas.GeoDataFrame({"geometry": geo_left})
+        right = geopandas.GeoDataFrame({"geometry": geo_right})
+        expected_gdf = left.iloc[expected_left].copy()
+        expected_gdf["index_right"] = expected_right
+        joined = sjoin(left, right, how=how, predicate="dwithin", distance=distance)
+        assert_frame_equal(expected_gdf.sort_index(), joined.sort_index())
+
 
 @pytest.mark.usefixtures("_setup_class_nybb_filename")
 class TestSpatialJoinNYBB: