diff --git a/python/dask_cudf/dask_cudf/backends.py b/python/dask_cudf/dask_cudf/backends.py index 1c66f256d19..7b8bd5c3f87 100644 --- a/python/dask_cudf/dask_cudf/backends.py +++ b/python/dask_cudf/dask_cudf/backends.py @@ -18,7 +18,6 @@ _scalar_from_dtype, is_arraylike, is_scalar, - make_meta, ) try: @@ -26,6 +25,13 @@ except ImportError: from dask.dataframe.utils import make_meta as make_meta_obj +try: + from dask.dataframe.dispatch import ( + make_meta_dispatch as make_meta_dispatch, + ) +except ImportError: + from dask.dataframe.utils import make_meta as make_meta_dispatch + import cudf from cudf.utils.dtypes import is_string_dtype @@ -121,12 +127,12 @@ def meta_nonempty_cudf(x): return res -@make_meta.register((cudf.Series, cudf.DataFrame)) +@make_meta_dispatch.register((cudf.Series, cudf.DataFrame)) def make_meta_cudf(x, index=None): return x.head(0) -@make_meta.register(cudf.Index) +@make_meta_dispatch.register(cudf.Index) def make_meta_cudf_index(x, index=None): return x[:0] @@ -173,7 +179,7 @@ def make_meta_object_cudf(x, index=None): return x[:0] if index is not None: - index = make_meta(index) + index = make_meta_dispatch(index) if isinstance(x, dict): return cudf.DataFrame( diff --git a/python/dask_cudf/dask_cudf/core.py b/python/dask_cudf/dask_cudf/core.py index dc63c5c435e..8af6188e625 100644 --- a/python/dask_cudf/dask_cudf/core.py +++ b/python/dask_cudf/dask_cudf/core.py @@ -14,7 +14,13 @@ from dask.compatibility import apply from dask.context import _globals from dask.core import flatten -from dask.dataframe.core import Scalar, finalize, handle_out, map_partitions +from dask.dataframe.core import ( + Scalar, + finalize, + handle_out, + make_meta as dask_make_meta, + map_partitions, +) from dask.dataframe.utils import raise_on_meta_error from dask.highlevelgraph import HighLevelGraph from dask.optimization import cull, fuse @@ -26,11 +32,6 @@ from dask_cudf import sorting from dask_cudf.accessors import ListMethods -try: - from dask.dataframe.utils import make_meta_util as dask_make_meta -except ImportError: - from dask.dataframe.core import make_meta as dask_make_meta - DASK_VERSION = LooseVersion(dask.__version__) diff --git a/python/dask_cudf/dask_cudf/tests/test_core.py b/python/dask_cudf/dask_cudf/tests/test_core.py index e7dff10b527..2f73534b45a 100644 --- a/python/dask_cudf/dask_cudf/tests/test_core.py +++ b/python/dask_cudf/dask_cudf/tests/test_core.py @@ -9,13 +9,7 @@ import dask from dask import dataframe as dd -from dask.dataframe.core import meta_nonempty - -try: - from dask.dataframe.utils import make_meta_util as dask_make_meta -except ImportError: - from dask.dataframe.core import make_meta as dask_make_meta - +from dask.dataframe.core import make_meta as dask_make_meta, meta_nonempty from dask.utils import M import cudf @@ -827,9 +821,9 @@ def test_merging_categorical_columns(): def test_correct_meta(): try: - from dask.dataframe.utils import make_meta_util # noqa: F401 + from dask.dataframe.dispatch import make_meta_obj # noqa: F401 except ImportError: - pytest.skip("need make_meta_util to be preset") + pytest.skip("need make_meta_obj to be preset") # Need these local imports in this specific order. # For context: https://github.com/rapidsai/cudf/issues/7946