Repository: dask/dask

Issue:
gpuCI failing
Yesterday out gpuCI builds started failing

```
16:28:38 FAILED dask/dataframe/tests/test_shuffle.py::test_set_index_interpolate[cudf]
16:28:38 FAILED dask/dataframe/tests/test_shuffle.py::test_set_index_interpolate_int[cudf]
16:28:38 FAILED dask/dataframe/tests/test_shuffle.py::test_set_index_interpolate_large_uint[cudf]
```

with error like `TypeError: No dispatch for <class 'cudf.core.series.Series'>` (see [this build](https://gpuci.gpuopenanalytics.com/job/dask/job/dask/job/prb/job/dask-prb/211/CUDA_VER=11.2,LINUX_VER=ubuntu18.04,PYTHON_VER=3.8,RAPIDS_VER=21.10/console) for an example). I believe this is because we were waiting on https://github.com/rapidsai/cudf/pull/9031 to be merged. That PR has since been merged, but gpuCI is still failing, now with a different error (see [this build](https://gpuci.gpuopenanalytics.com/job/dask/job/dask/job/prb/job/dask-prb/219/CUDA_VER=11.2,LINUX_VER=ubuntu18.04,PYTHON_VER=3.8,RAPIDS_VER=21.10/console) for an example):

```python
11:29:33 _______________________ test_set_index_interpolate[cudf] _______________________
11:29:33 [gw3] linux -- Python 3.8.10 /opt/conda/envs/dask/bin/python
11:29:33 
11:29:33 engine = 'cudf'
11:29:33 
11:29:33     @pytest.mark.gpu
11:29:33     @pytest.mark.parametrize("engine", ["pandas", "cudf"])
11:29:33     def test_set_index_interpolate(engine):
11:29:33         if engine == "cudf":
11:29:33             # NOTE: engine == "cudf" requires cudf/dask_cudf,
11:29:33             # will be skipped by non-GPU CI.
11:29:33     
11:29:33             cudf = pytest.importorskip("cudf")
11:29:33             dask_cudf = pytest.importorskip("dask_cudf")
11:29:33     
11:29:33         df = pd.DataFrame({"x": [4, 1, 1, 3, 3], "y": [1.0, 1, 1, 1, 2]})
11:29:33     
11:29:33         if engine == "cudf":
11:29:33             gdf = cudf.from_pandas(df)
11:29:33             d = dask_cudf.from_cudf(gdf, npartitions=3)
11:29:33         else:
11:29:33             d = dd.from_pandas(df, 2)
11:29:33     
11:29:33         d1 = d.set_index("x", npartitions=3)
11:29:33         assert d1.npartitions == 3
11:29:33         if engine == "cudf":
11:29:33             # cuDF has to use "linear" interpolation as "nearest" interpolation
11:29:33             # is not supported by CuPy. The result for d1.divisions is different
11:29:33             # in this case.
11:29:33 >           assert set(d1.divisions) == set([1, 2, 3, 4])
11:29:33 E           assert {1, 2, 4} == {1, 2, 3, 4}
11:29:33 E             Extra items in the right set:
11:29:33 E             3
11:29:33 E             Full diff:
11:29:33 E             - {1, 2, 3, 4}
11:29:33 E             ?        ---
11:29:33 E             + {1, 2, 4}
11:29:33 
11:29:33 dask/dataframe/tests/test_shuffle.py:647: AssertionError
```

cc @dask/gpu 


Discussion:

Resolving patch:
diff --git a/dask/array/percentile.py b/dask/array/percentile.py
--- a/dask/array/percentile.py
+++ b/dask/array/percentile.py
@@ -12,8 +12,6 @@
 
 @wraps(np.percentile)
 def _percentile(a, q, interpolation="linear"):
-    import pandas as pd
-
     n = len(a)
     if not len(a):
         return None, n
@@ -21,15 +19,17 @@ def _percentile(a, q, interpolation="linear"):
         q = list(q)
     if a.dtype.name == "category":
         result = np.percentile(a.cat.codes, q, interpolation=interpolation)
+        import pandas as pd
 
         return pd.Categorical.from_codes(result, a.dtype.categories, a.dtype.ordered), n
-    if isinstance(a.dtype, pd.core.dtypes.dtypes.DatetimeTZDtype) or np.issubdtype(
-        a.dtype, np.datetime64
-    ):
+    if type(a.dtype).__name__ == "DatetimeTZDtype":
+        import pandas as pd
+
         if isinstance(a, (pd.Series, pd.Index)):
-            values = a.values
-        else:
-            values = a
+            a = a.values
+
+    if np.issubdtype(a.dtype, np.datetime64):
+        values = a
         a2 = values.view("i8")
         result = np.percentile(a2, q, interpolation=interpolation).astype(values.dtype)
         if q[0] == 0:
diff --git a/dask/dataframe/backends.py b/dask/dataframe/backends.py
--- a/dask/dataframe/backends.py
+++ b/dask/dataframe/backends.py
@@ -535,6 +535,11 @@ def is_categorical_dtype_pandas(obj):
     return pd.api.types.is_categorical_dtype(obj)
 
 
+@percentile_dispatch.register((pd.Series, np.ndarray, pd.Index))
+def percentile(a, q, interpolation="linear"):
+    return _percentile(a, q, interpolation)
+
+
 ######################################
 # cuDF: Pandas Dataframes on the GPU #
 ######################################
@@ -547,10 +552,6 @@ def is_categorical_dtype_pandas(obj):
 @meta_nonempty.register_lazy("cudf")
 @make_meta_dispatch.register_lazy("cudf")
 @make_meta_obj.register_lazy("cudf")
+@percentile_dispatch.register_lazy("cudf")
 def _register_cudf():
     import dask_cudf  # noqa: F401
-
-
-@percentile_dispatch.register((pd.Series, np.ndarray, pd.Index))
-def percentile(a, q, interpolation="linear"):
-    return _percentile(a, q, interpolation)
diff --git a/dask/dataframe/dispatch.py b/dask/dataframe/dispatch.py
--- a/dask/dataframe/dispatch.py
+++ b/dask/dataframe/dispatch.py
@@ -26,8 +26,7 @@
 
 
 def _percentile(a, q, interpolation="linear"):
-    func = percentile_dispatch.dispatch(type(a))
-    return func(a, q, interpolation)
+    return percentile_dispatch(a, q, interpolation)
 
 
 def concat(


Test patch:
diff --git a/dask/dataframe/tests/test_shuffle.py b/dask/dataframe/tests/test_shuffle.py
--- a/dask/dataframe/tests/test_shuffle.py
+++ b/dask/dataframe/tests/test_shuffle.py
@@ -640,13 +640,7 @@ def test_set_index_interpolate(engine):
 
     d1 = d.set_index("x", npartitions=3)
     assert d1.npartitions == 3
-    if engine == "cudf":
-        # cuDF has to use "linear" interpolation as "nearest" interpolation
-        # is not supported by CuPy. The result for d1.divisions is different
-        # in this case.
-        assert set(d1.divisions) == set([1, 2, 3, 4])
-    else:
-        assert set(d1.divisions) == set([1, 2, 4])
+    assert set(d1.divisions) == set([1, 2, 4])
 
     d2 = d.set_index("y", npartitions=3)
     assert d2.divisions[0] == 1.0