diff --git a/docs/setup/release-notes.md b/docs/setup/release-notes.md index 0b24653fab9..c1d7641ab26 100644 --- a/docs/setup/release-notes.md +++ b/docs/setup/release-notes.md @@ -65,6 +65,7 @@ * [GH-3273] - Implement distributed GeoSeries and GeoDataFrame polygonal coverage validation and invalid-edge diagnostics * [GH-3281] - Implement `GeoDataFrame.from_features` for in-memory GeoJSON-like features * [GH-3288] - Implement `GeoDataFrame.from_dict` for in-memory dictionaries +* [GH-3344] - Expose `SpatialIndex.valid_query_predicates` ### Bug Fixes diff --git a/python/sedona/spark/geopandas/sindex.py b/python/sedona/spark/geopandas/sindex.py index 7ad592ba952..78bebda5d4c 100644 --- a/python/sedona/spark/geopandas/sindex.py +++ b/python/sedona/spark/geopandas/sindex.py @@ -80,6 +80,21 @@ def __init__(self, geometry, index_type="strtree", column_name=None): "Invalid type for `geometry`. Expected np.array, GeoSeries, or PySparkDataFrame." ) + @property + def valid_query_predicates(self) -> set: + """ + Return the supported values for the ``query`` predicate. + + .. versionadded:: 2.0.0 + + Returns + ------- + set + A new set containing ``None``, ``"intersects"``, and ``"contains"``. + ``None`` selects the default ``"intersects"`` behavior. + """ + return {None, *ALLOWED_PREDICATES} + def query(self, geometry: BaseGeometry, predicate: str = None, sort: bool = False): """ Query the spatial index for geometries that intersect the given geometry. diff --git a/python/tests/geopandas/test_sindex.py b/python/tests/geopandas/test_sindex.py index 6e68fc40b1a..8ad8937ce4f 100644 --- a/python/tests/geopandas/test_sindex.py +++ b/python/tests/geopandas/test_sindex.py @@ -94,6 +94,32 @@ def test_geoseries_sindex_property_exists(self): assert hasattr(self.polygons, "sindex") assert hasattr(self.lines, "sindex") + @pytest.mark.parametrize("distributed", [False, True]) + def test_valid_query_predicates(self, distributed): + geometries = [Point(0, 0), Point(1, 1), Point(2, 2)] + sindex = ( + GeoSeries(geometries).sindex + if distributed + else SpatialIndex(np.array(geometries)) + ) + assert sindex.valid_query_predicates == {None, "intersects", "contains"} + + for predicate in sindex.valid_query_predicates: + result = sindex.query(box(0.5, 0.5, 1.5, 1.5), predicate=predicate) + assert list(result) == ([Point(1, 1)] if distributed else [1]) + + @pytest.mark.parametrize("geometries", [[], [Point(0, 0)]]) + def test_valid_query_predicates_returns_independent_set(self, geometries): + sindex = SpatialIndex(np.array(geometries, dtype=object)) + predicates = sindex.valid_query_predicates + predicates.clear() + predicates.add("within") + + assert sindex.valid_query_predicates == {None, "intersects", "contains"} + if geometries: + with pytest.raises(ValueError, match="Predicate must be"): + sindex.query(Point(0, 0), predicate="within") + def test_geodataframe_sindex_property_exists(self): """Test that the sindex property exists on GeoDataFrame.""" assert hasattr(self.points.to_geoframe(), "sindex")