Add apply_boolean_mask Feature

python/cuspatial/cuspatial/core/geodataframe.py

@@ -181,6 +181,18 @@ def _slice(self: T, arg: slice) -> T:
         )
         return self.__class__(result)
 
+    def _apply_boolean_mask(self, mask) -> T:
+        geo_columns, data_columns = self._split_out_geometry_columns()
+        data = data_columns._apply_boolean_mask(mask)
+
+        geo = GeoDataFrame(
+            {name: geo_columns[name][mask] for name in geo_columns}
+        )
+
+        res = self.__class__._from_data(self._recombine_columns(geo, data))
+        res.index = data.index
+        return res
+
     def _gather(
         self, gather_map, keep_index=True, nullify=False, check_bounds=True
     ):

python/cuspatial/cuspatial/tests/conftest.py

@@ -286,3 +286,22 @@ def slice_twenty():
         slice(12, 16),
         slice(16, 20),
     ]
+
+
+@pytest.fixture(params=["MaskOdd", "MaskEven", "MaskNone", "MaskAll"])
+def mask_factory(request):
+    kind = request.param
+
+    def factory(length):
+        mask = pd.Series([False] * length)
+        if kind == "MaskOdd":
+            mask[0::2] = True
+        elif kind == "MaskEven":
+            mask[1::2] = True
+        elif kind == "MaskNone":
+            pass
+        elif kind == "MaskAll":
+            mask[:] = True
+        return mask
+
+    return factory

python/cuspatial/cuspatial/tests/test_geodataframe.py

@@ -3,6 +3,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+from geopandas.testing import assert_geodataframe_equal
 from shapely.affinity import rotate
 from shapely.geometry import (
     LineString,
@@ -446,3 +447,25 @@ def test_cudf_dataframe_init():
     df = cudf.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
     gdf = cuspatial.GeoDataFrame(df)
     assert_eq_geo_df(gdf.to_pandas(), df.to_pandas())
+
+
+def test_apply_boolean_mask(gpdf, mask_factory):
+    mask = mask_factory(len(gpdf))
+
+    expected = gpdf[mask]
+
+    d_gpdf = cuspatial.from_geopandas(gpdf)
+    got = d_gpdf[mask]
+
+    assert_geodataframe_equal(expected, got.to_geopandas())
+
+
+def test_apply_boolean_mask_length_one(mask_factory):
+    geodf = gpd.GeoDataFrame({"geometry": [Point(0, 0)]})
+    mask = mask_factory(len(geodf))
+    expected = geodf[mask]
+
+    d_geodf = cuspatial.from_geopandas(geodf)
+    got = d_geodf[mask]
+
+    assert_geodataframe_equal(expected, got.to_geopandas())