Update to support cudf 2024.06 (#1344)

Co-authored-by: James A. Bednar <[email protected]>

datashader/tests/test_dask.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
-
 import dask.dataframe as dd
 import numpy as np
 import pandas as pd
 import xarray as xr
 
 from dask.context import config
 from numpy import nan
+from packaging.version import Version
 
 import datashader as ds
 from datashader.datatypes import RaggedArray
@@ -42,12 +42,24 @@ def _dask_expr():
 def _dask_cudf():
     import dask_cudf
     _dask = dd.from_pandas(_pandas(), npartitions=2)
-    return dask_cudf.from_dask_dataframe(_dask)
+    if Version(dask_cudf.__version__) >= Version("24.06"):
+        return _dask.to_backend("cudf")
+    else:
+        return dask_cudf.from_dask_dataframe(_dask)
+
+@dask_switcher(query=True, extras=["dask_cudf"])
+def _dask_expr_cudf():
+    import dask_cudf
+    if Version(dask_cudf.__version__) < Version("24.06"):
+        pytest.skip("dask-expr requires dask-cudf 24.06 or later")
+    _dask = dd.from_pandas(_pandas(), npartitions=2)
+    return _dask.to_backend("cudf")
 
 _backends = [
     pytest.param(_dask, id="dask"),
     pytest.param(_dask_expr, id="dask-expr"),
     pytest.param(_dask_cudf, marks=pytest.mark.gpu, id="dask-cudf"),
+    pytest.param(_dask_expr_cudf, marks=pytest.mark.gpu, id="dask-expr-cudf"),
 ]
 
 @pytest.fixture(params=_backends)
@@ -83,16 +95,38 @@ def _dask_cudf_DataFrame(*args, **kwargs):
     import cudf
     import dask_cudf
     if kwargs.pop("geo", False):
+        # As of dask-cudf version 24.06, dask-cudf is not
+        # compatible with spatialpandas version 0.4.10
+        pytest.skip("dask-cudf currently does not work with spatialpandas")
+    cdf = cudf.DataFrame.from_pandas(
+        pd.DataFrame(*args, **kwargs), nan_as_null=False
+    )
+    return dask_cudf.from_cudf(cdf, npartitions=2)
+
+
+@dask_switcher(query=True, extras=["dask_cudf"])
+def _dask_expr_cudf_DataFrame(*args, **kwargs):
+    import cudf
+    import dask_cudf
+
+    if Version(dask_cudf.__version__) < Version("24.06"):
+        pytest.skip("dask-expr requires dask-cudf 24.06 or later")
+
+    if kwargs.pop("geo", False):
+        # As of dask-cudf version 24.06, dask-cudf is not
+        # compatible with spatialpandas version 0.4.10
         pytest.skip("dask-cudf currently does not work with spatialpandas")
     cdf = cudf.DataFrame.from_pandas(
         pd.DataFrame(*args, **kwargs), nan_as_null=False
     )
     return dask_cudf.from_cudf(cdf, npartitions=2)
 
+
 _backends = [
     pytest.param(_dask_DataFrame, id="dask"),
     pytest.param(_dask_expr_DataFrame, id="dask-expr"),
     pytest.param(_dask_cudf_DataFrame, marks=pytest.mark.gpu, id="dask-cudf"),
+    pytest.param(_dask_expr_cudf_DataFrame, marks=pytest.mark.gpu, id="dask-expr-cudf"),
 ]
 
 @pytest.fixture(params=_backends)
@@ -127,16 +161,20 @@ def floats(n):
 @pytest.mark.gpu
 def test_check_query_setting():
     import os
-    from subprocess import check_output
+    from subprocess import check_output, SubprocessError
 
     # dask-cudf does not support query planning as of 24.04.
     # So we check that it is not set outside of Python.
     assert os.environ.get('DASK_DATAFRAME__QUERY_PLANNING', 'false').lower() != 'true'
 
     # This also have problem with the global setting so we check
-    cmd = ['dask', 'config', 'get', 'dataframe.query-planning']
-    output = check_output(cmd, text=True).strip().lower()
-    assert output != 'true'
+    try:
+        cmd = ['dask', 'config', 'get', 'dataframe.query-planning']
+        output = check_output(cmd, text=True).strip().lower()
+        assert output != 'true'
+    except SubprocessError:
+        # Newer version will error out if not set
+        pass
 
 
 def test_count(ddf, npartitions):
@@ -1578,7 +1616,7 @@ def test_auto_range_line(DataFrame):
     }, dtype='Ragged[float32]'), dict(x='x', y='y', axis=1))
 ])
 def test_area_to_zero_fixedrange(DataFrame, df_kwargs, cvs_kwargs):
-    if DataFrame is _dask_cudf_DataFrame:
+    if DataFrame in (_dask_cudf_DataFrame, _dask_expr_cudf_DataFrame):
         if df_kwargs.get('dtype', '').startswith('Ragged'):
             pytest.skip("Ragged array not supported with cudf")
 
@@ -1670,7 +1708,7 @@ def test_area_to_zero_fixedrange(DataFrame, df_kwargs, cvs_kwargs):
     }, dtype='Ragged[float32]'), dict(x='x', y='y', axis=1))
 ])
 def test_area_to_zero_autorange(DataFrame, df_kwargs, cvs_kwargs):
-    if DataFrame is _dask_cudf_DataFrame:
+    if DataFrame in (_dask_cudf_DataFrame, _dask_expr_cudf_DataFrame):
         if df_kwargs.get('dtype', '').startswith('Ragged'):
             pytest.skip("Ragged array not supported with cudf")
 
@@ -1747,7 +1785,7 @@ def test_area_to_zero_autorange(DataFrame, df_kwargs, cvs_kwargs):
     }, dtype='Ragged[float32]'), dict(x='x', y='y', axis=1))
 ])
 def test_area_to_zero_autorange_gap(DataFrame, df_kwargs, cvs_kwargs):
-    if DataFrame is _dask_cudf_DataFrame:
+    if DataFrame in (_dask_cudf_DataFrame, _dask_expr_cudf_DataFrame):
         if df_kwargs.get('dtype', '').startswith('Ragged'):
             pytest.skip("Ragged array not supported with cudf")
 
@@ -1850,7 +1888,7 @@ def test_area_to_zero_autorange_gap(DataFrame, df_kwargs, cvs_kwargs):
     }, dtype='Ragged[float32]'), dict(x='x', y='y', y_stack='y_stack', axis=1))
 ])
 def test_area_to_line_autorange(DataFrame, df_kwargs, cvs_kwargs):
-    if DataFrame is _dask_cudf_DataFrame:
+    if DataFrame in (_dask_cudf_DataFrame, _dask_expr_cudf_DataFrame):
         if df_kwargs.get('dtype', '').startswith('Ragged'):
             pytest.skip("Ragged array not supported with cudf")
 
@@ -1937,7 +1975,7 @@ def test_area_to_line_autorange(DataFrame, df_kwargs, cvs_kwargs):
     }, dtype='Ragged[float32]'), dict(x='x', y='y', y_stack='y_stack', axis=1))
 ])
 def test_area_to_line_autorange_gap(DataFrame, df_kwargs, cvs_kwargs):
-    if DataFrame is _dask_cudf_DataFrame:
+    if DataFrame in (_dask_cudf_DataFrame, _dask_expr_cudf_DataFrame):
         if df_kwargs.get('dtype', '').startswith('Ragged'):
             pytest.skip("Ragged array not supported with cudf")
 

setup.py

@@ -21,6 +21,7 @@
     'requests',
     'scipy',
     'toolz',
+    'packaging',
     'xarray',
 ]