From ea7522ac8e371e386e72c45edd11a46447f9c055 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Simon=20H=C3=B8xbro=20Hansen?= <simon.hansen@me.com>
Date: Mon, 6 May 2024 20:30:33 +0200
Subject: [PATCH] Update cuspatial

---
 holoviews/element/selection.py            | 46 ++++++++++++++++++-----
 holoviews/tests/element/test_selection.py |  1 -
 2 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/holoviews/element/selection.py b/holoviews/element/selection.py
index ba5fbbd285..a6c8ff43ac 100644
--- a/holoviews/element/selection.py
+++ b/holoviews/element/selection.py
@@ -80,6 +80,38 @@ def spatial_select_gridded(xvals, yvals, geometry):
         sel_mask = spatial_select_columnar(xvals.flatten(), yvals.flatten(), geometry)
         return sel_mask.reshape(xvals.shape)
 
+def _cuspatial_old(xvals, yvals, geometry):
+    import cudf
+    import cuspatial
+
+    result = cuspatial.point_in_polygon(
+        xvals,
+        yvals,
+        cudf.Series([0], index=["selection"]),
+        [0],
+        geometry[:, 0],
+        geometry[:, 1],
+    )
+    return result.values
+
+
+def _cuspatial_new(xvals, yvals, geometry):
+    import cudf
+    import cuspatial
+    import geopandas
+    from shapely.geometry import Polygon
+
+    df = cudf.DataFrame({'x':xvals, 'y':yvals})
+    points = cuspatial.GeoSeries.from_points_xy(
+       df.interleave_columns().astype('float')
+    )
+    polygons = cuspatial.GeoSeries(
+       geopandas.GeoSeries(Polygon(geometry)), index=["selection"]
+    )
+    result = cuspatial.point_in_polygon(points,polygons)
+    return result.values.ravel()
+
+
 def spatial_select_columnar(xvals, yvals, geometry, geom_method=None):
     if 'cudf' in sys.modules:
         import cudf
@@ -88,16 +120,10 @@ def spatial_select_columnar(xvals, yvals, geometry, geom_method=None):
             xvals = xvals.values.astype('float')
             yvals = yvals.values.astype('float')
             try:
-                import cuspatial
-                result = cuspatial.point_in_polygon(  # TODO: now only takes two arguments
-                    xvals,
-                    yvals,
-                    cudf.Series([0], index=["selection"]),
-                    [0],
-                    geometry[:, 0],
-                    geometry[:, 1],
-                )
-                return result.values
+                try:
+                    return _cuspatial_old(xvals, yvals, geometry)
+                except TypeError:
+                    return _cuspatial_new(xvals, yvals, geometry)
             except ImportError:
                 xvals = cp.asnumpy(xvals)
                 yvals = cp.asnumpy(yvals)
diff --git a/holoviews/tests/element/test_selection.py b/holoviews/tests/element/test_selection.py
index 7562e72ef9..517653e8e4 100644
--- a/holoviews/tests/element/test_selection.py
+++ b/holoviews/tests/element/test_selection.py
@@ -692,7 +692,6 @@ def test_cudf(self, geometry, pt_mask, pandas_df, _method, unimport):
             mask = spatial_select_columnar(df.x, df.y, geometry, _method)
             assert np.array_equal(cp.asnumpy(mask), pt_mask)
 
-        @pytest.mark.xfail(reason='cuspatial.point_in_polygon API has changed')
         @pytest.mark.gpu
         def test_cuspatial(self, geometry, pt_mask, pandas_df, _method):
             import cudf