Enable dask dispatch to cuDF's is_categorical_dtype for cuDF objects

conda/environments/cudf_dev_cuda10.1.yml

@@ -43,7 +43,7 @@ dependencies:
   - mypy=0.782
   - typing_extensions
   - pre_commit
-  - dask>=2.22.0
+  - dask>=2021.3.1
   - distributed>=2.22.0
   - streamz
   - dlpack

conda/environments/cudf_dev_cuda10.2.yml

@@ -43,7 +43,7 @@ dependencies:
   - mypy=0.782
   - typing_extensions
   - pre_commit
-  - dask>=2.22.0
+  - dask>=2021.3.1
   - distributed>=2.22.0
   - streamz
   - dlpack

conda/environments/cudf_dev_cuda11.0.yml

@@ -43,7 +43,7 @@ dependencies:
   - mypy=0.782
   - typing_extensions
   - pre_commit
-  - dask>=2.22.0
+  - dask>=2021.3.1
   - distributed>=2.22.0
   - streamz
   - dlpack

conda/recipes/dask-cudf/meta.yaml

@@ -23,12 +23,12 @@ requirements:
   host:
     - python
     - cudf {{ version }}=*_{{ GIT_DESCRIBE_NUMBER }}
-    - dask >=2.22.0
+    - dask>=2021.3.1
     - distributed >=2.22.0
   run:
     - python
     - cudf {{ version }}=*_{{ GIT_DESCRIBE_NUMBER }}
-    - dask >=2.22.0
+    - dask>=2021.3.1
     - distributed >=2.22.0
 
 test:

python/dask_cudf/dask_cudf/backends.py

@@ -6,7 +6,11 @@
 
 from dask.dataframe.categorical import categorical_dtype_dispatch
 from dask.dataframe.core import get_parallel_type, make_meta, meta_nonempty
-from dask.dataframe.methods import concat_dispatch, tolist_dispatch
+from dask.dataframe.methods import (
+    concat_dispatch,
+    is_categorical_dtype_dispatch,
+    tolist_dispatch,
+)
 from dask.dataframe.utils import (
     UNKNOWN_CATEGORIES,
     _nonempty_scalar,
@@ -220,6 +224,13 @@ def tolist_cudf(obj):
     return obj.to_arrow().to_pylist()
 
 
+@is_categorical_dtype_dispatch.register(
+    (cudf.Series, cudf.Index, cudf.CategoricalDtype, Series)
+)
+def is_categorical_dtype_cudf(obj):
+    return cudf.utils.dtypes.is_categorical_dtype(obj)
+
+
 try:
 
     from dask.dataframe.utils import group_split_dispatch, hash_object_dispatch

python/dask_cudf/dask_cudf/tests/test_dispatch.py

@@ -0,0 +1,16 @@
+import pandas as pd
+
+from dask.dataframe.methods import is_categorical_dtype
+
+import cudf
+
+
+def test_is_categorical_dispatch():
+    assert is_categorical_dtype(pd.CategoricalDtype([1, 2, 3]))
+    assert is_categorical_dtype(cudf.CategoricalDtype([1, 2, 3]))
+
+    assert is_categorical_dtype(cudf.Series([1, 2, 3], dtype="category"))
+    assert is_categorical_dtype(pd.Series([1, 2, 3], dtype="category"))
+
+    assert is_categorical_dtype(pd.Index([1, 2, 3], dtype="category"))
+    assert is_categorical_dtype(cudf.Index([1, 2, 3], dtype="category"))

python/dask_cudf/dask_cudf/tests/test_onehot.py

@@ -3,10 +3,10 @@
 
 from dask import dataframe as dd
 
-import dask_cudf
-
 import cudf
 
+import dask_cudf
+
 
 def test_get_dummies_cat():
     df = pd.DataFrame({"C": [], "A": []})
@@ -101,3 +101,22 @@ def test_get_dummies_large():
         dd.get_dummies(gddf).compute(),
         check_dtype=False,
     )
+
+
+def test_get_dummies_categorical():
+    # https://github.com/rapidsai/cudf/issues/7111
+    gdf = cudf.DataFrame({"A": ["a", "b", "b"], "B": [1, 2, 3]})
+    pdf = gdf.to_pandas()
+
+    gddf = dask_cudf.from_cudf(gdf, npartitions=1)
+    gddf = gddf.categorize(columns=["B"])
+
+    pddf = dd.from_pandas(pdf, npartitions=1)
+    pddf = pddf.categorize(columns=["B"])
+
+    expect = dd.get_dummies(pddf, columns=["B"])
+    got = dd.get_dummies(gddf, columns=["B"])
+
+    dd.assert_eq(
+        expect, got,
+    )