Introduce basic "cudf" backend for Dask Expressions

ci/test_wheel_dask_cudf.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
 
 set -eou pipefail
 
@@ -28,3 +28,7 @@ python -m pip install $(echo ./dist/dask_cudf*.whl)[test]
 
 # Run tests in dask_cudf/tests and dask_cudf/io/tests
 python -m pytest -n 8 ./python/dask_cudf/dask_cudf/
+
+# Run tests in dask_cudf/tests and dask_cudf/io/tests with dask-expr
+echo "Running dask-cudf tests with dask-expr enabled..."
+DASK_DATAFRAME__QUERY_PLANNING=True python -m pytest -n 8 ./python/dask_cudf/dask_cudf/

conda/environments/all_cuda-118_arch-x86_64.yaml

@@ -102,5 +102,6 @@ dependencies:
 - typing_extensions>=4.0.0
 - zlib>=1.2.13
 - pip:
+  - git+https://github.com/dask-contrib/dask-expr.git@b588a9e15e90e0567061664ffc01374786686e20
   - git+https://github.com/python-streamz/streamz.git@master
 name: all_cuda-118_arch-x86_64

conda/environments/all_cuda-120_arch-x86_64.yaml

@@ -100,5 +100,6 @@ dependencies:
 - typing_extensions>=4.0.0
 - zlib>=1.2.13
 - pip:
+  - git+https://github.com/dask-contrib/dask-expr.git@b588a9e15e90e0567061664ffc01374786686e20
   - git+https://github.com/python-streamz/streamz.git@master
 name: all_cuda-120_arch-x86_64

dependencies.yaml

@@ -661,6 +661,18 @@ dependencies:
         packages:
           - dask-cuda==24.4.*
           - *numba
+      - output_types: conda
+        packages:
+          - pip
+          - pip:
+              # This should eventually move to rapids-dask-dependency
+              - &dask_expr_tip git+https://github.com/dask-contrib/dask-expr.git@b588a9e15e90e0567061664ffc01374786686e20
+      - output_types: requirements
+        packages:
+          - *dask_expr_tip
+      - output_types: pyproject
+        packages:
+          - dask-expr@git+https://github.com/dask-contrib/dask-expr.git@b588a9e15e90e0567061664ffc01374786686e20
   depends_on_cudf:
     common:
       - output_types: conda

python/dask_cudf/dask_cudf/__init__.py

@@ -1,29 +1,69 @@
-# Copyright (c) 2018-2023, NVIDIA CORPORATION.
+# Copyright (c) 2018-2024, NVIDIA CORPORATION.
 
+import dask.dataframe as dd
+from dask import config
 from dask.dataframe import from_delayed
 
 import cudf
 
 from . import backends
 from ._version import __git_commit__, __version__
-from .core import DataFrame, Series, concat, from_cudf, from_dask_dataframe
-from .groupby import groupby_agg
-from .io import read_csv, read_json, read_orc, read_text, to_orc
+from .core import concat, from_cudf, from_dask_dataframe
+from .expr import DASK_EXPR_ENABLED
+
+
+def read_csv(*args, **kwargs):
+    with config.set({"dataframe.backend": "cudf"}):
+        return dd.read_csv(*args, **kwargs)
+
+
+def read_json(*args, **kwargs):
+    with config.set({"dataframe.backend": "cudf"}):
+        return dd.read_json(*args, **kwargs)
+
+
+def read_orc(*args, **kwargs):
+    with config.set({"dataframe.backend": "cudf"}):
+        return dd.read_orc(*args, **kwargs)
+
+
+def read_parquet(*args, **kwargs):
+    with config.set({"dataframe.backend": "cudf"}):
+        return dd.read_parquet(*args, **kwargs)
+
+
+def raise_not_implemented_error(attr_name):
+    def inner_func(*args, **kwargs):
+        raise NotImplementedError(
+            f"Top-level {attr_name} API is not available for dask-expr."
+        )
+
+    return inner_func
+
+
+if DASK_EXPR_ENABLED:
+    from .expr._collection import DataFrame, Index, Series
+
+    groupby_agg = raise_not_implemented_error("groupby_agg")
+    read_text = raise_not_implemented_error("read_text")
+    to_orc = raise_not_implemented_error("to_orc")
+else:
+    from .core import DataFrame, Index, Series
+    from .groupby import groupby_agg
+    from .io import read_text, to_orc
 
-try:
-    from .io import read_parquet
-except ImportError:
-    pass
 
 __all__ = [
     "DataFrame",
     "Series",
+    "Index",
     "from_cudf",
     "from_dask_dataframe",
     "concat",
     "from_delayed",
 ]
 
+
 if not hasattr(cudf.DataFrame, "mean"):
     cudf.DataFrame.mean = None
 del cudf

python/dask_cudf/dask_cudf/backends.py

@@ -625,13 +625,68 @@ def read_csv(*args, **kwargs):
 
     @staticmethod
     def read_hdf(*args, **kwargs):
-        from dask_cudf import from_dask_dataframe
-
         # HDF5 reader not yet implemented in cudf
         warnings.warn(
             "read_hdf is not yet implemented in cudf/dask_cudf. "
             "Moving to cudf from pandas. Expect poor performance!"
         )
-        return from_dask_dataframe(
-            _default_backend(dd.read_hdf, *args, **kwargs)
+        return _default_backend(dd.read_hdf, *args, **kwargs).to_backend(
+            "cudf"
+        )
+
+
+# Define "cudf" backend entrypoint for dask-expr
+class CudfDXBackendEntrypoint(DataFrameBackendEntrypoint):
+    """Backend-entrypoint class for Dask-Expressions
+
+    This class is registered under the name "cudf" for the
+    ``dask-expr.dataframe.backends`` entrypoint in ``setup.cfg``.
+    Dask-DataFrame will use the methods defined in this class
+    in place of ``dask_expr.<creation-method>`` when the
+    "dataframe.backend" configuration is set to "cudf":
+
+    Examples
+    --------
+    >>> import dask
+    >>> import dask_expr
+    >>> with dask.config.set({"dataframe.backend": "cudf"}):
+    ...     ddf = dx.from_dict({"a": range(10)})
+    >>> type(ddf._meta)
+    <class 'cudf.core.dataframe.DataFrame'>
+    """
+
+    @classmethod
+    def to_backend_dispatch(cls):
+        return CudfBackendEntrypoint.to_backend_dispatch()
+
+    @classmethod
+    def to_backend(cls, *args, **kwargs):
+        return CudfBackendEntrypoint.to_backend(*args, **kwargs)
+
+    @staticmethod
+    def from_dict(
+        data,
+        npartitions,
+        orient="columns",
+        dtype=None,
+        columns=None,
+        constructor=cudf.DataFrame,
+    ):
+        import dask_expr as dx
+
+        return _default_backend(
+            dx.from_dict,
+            data,
+            npartitions=npartitions,
+            orient=orient,
+            dtype=dtype,
+            columns=columns,
+            constructor=constructor,
         )
+
+
+# Import/register cudf-specific classes for dask-expr
+try:
+    import dask_cudf.expr  # noqa: F401
+except ImportError:
+    pass

python/dask_cudf/dask_cudf/core.py

@@ -8,7 +8,7 @@
 import pandas as pd
 from tlz import partition_all
 
-from dask import dataframe as dd
+from dask import config, dataframe as dd
 from dask.base import normalize_token, tokenize
 from dask.dataframe.core import (
     Scalar,
@@ -690,13 +690,20 @@ def from_cudf(data, npartitions=None, chunksize=None, sort=True, name=None):
             "dask_cudf does not support MultiIndex Dataframes."
         )
 
-    name = name or ("from_cudf-" + tokenize(data, npartitions or chunksize))
+    # Dask-expr doesn't support the `name` argument
+    name = {}
+    if not config.get("dataframe.query-planning", False):
+        name = {
+            "name": name
+            or ("from_cudf-" + tokenize(data, npartitions or chunksize))
+        }
+
     return dd.from_pandas(
         data,
         npartitions=npartitions,
         chunksize=chunksize,
         sort=sort,
-        name=name,
+        **name,
     )
 
 
@@ -711,7 +718,9 @@ def from_cudf(data, npartitions=None, chunksize=None, sort=True, name=None):
         rather than pandas objects.\n
         """
     )
-    + textwrap.dedent(dd.from_pandas.__doc__)
+    # TODO: `dd.from_pandas.__doc__` is empty when
+    # `DASK_DATAFRAME__QUERY_PLANNING=True`
+    + textwrap.dedent(dd.from_pandas.__doc__ or "")
 )
 
 

python/dask_cudf/dask_cudf/expr/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+from dask import config
+
+DASK_EXPR_ENABLED = False
+if config.get("dataframe.query-planning", False):
+    # Make sure custom expressions and collections are defined
+    try:
+        import dask_cudf.expr._collection
+        import dask_cudf.expr._expr
+
+        DASK_EXPR_ENABLED = True
+    except ImportError:
+        # Dask Expressions not installed.
+        # Dask DataFrame should have already thrown an error
+        # before we got here.
+        pass

python/dask_cudf/dask_cudf/expr/_collection.py

@@ -0,0 +1,66 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+from dask_expr import (
+    DataFrame as DXDataFrame,
+    FrameBase,
+    Index as DXIndex,
+    Series as DXSeries,
+    get_collection_type,
+)
+
+from dask import config
+
+import cudf
+
+##
+## Custom collection classes
+##
+
+
+class DataFrame(DXDataFrame):
+    @classmethod
+    def from_dict(cls, *args, **kwargs):
+        with config.set({"dataframe.backend": "cudf"}):
+            return DXDataFrame.from_dict(*args, **kwargs)
+
+    def groupby(
+        self,
+        by,
+        group_keys=True,
+        sort=None,
+        observed=None,
+        dropna=None,
+        **kwargs,
+    ):
+        from dask_cudf.expr._groupby import GroupBy
+
+        if isinstance(by, FrameBase) and not isinstance(by, DXSeries):
+            raise ValueError(
+                f"`by` must be a column name or list of columns, got {by}."
+            )
+
+        return GroupBy(
+            self,
+            by,
+            group_keys=group_keys,
+            sort=sort,
+            observed=observed,
+            dropna=dropna,
+            **kwargs,
+        )
+
+
+class Series(DXSeries):
+    def groupby(self, by, **kwargs):
+        from dask_cudf.expr._groupby import SeriesGroupBy
+
+        return SeriesGroupBy(self, by, **kwargs)
+
+
+class Index(DXIndex):
+    pass  # Same as pandas (for now)
+
+
+get_collection_type.register(cudf.DataFrame, lambda _: DataFrame)
+get_collection_type.register(cudf.Series, lambda _: Series)
+get_collection_type.register(cudf.BaseIndex, lambda _: Index)

python/dask_cudf/dask_cudf/expr/_expr.py

@@ -0,0 +1,34 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+from dask_expr._cumulative import CumulativeBlockwise, TakeLast
+
+##
+## Custom expression patching
+##
+
+
+class PatchCumulativeBlockwise(CumulativeBlockwise):
+    @property
+    def _args(self) -> list:
+        return self.operands[:1]
+
+    @property
+    def _kwargs(self) -> dict:
+        # Must pass axis and skipna as kwargs in cudf
+        return {"axis": self.axis, "skipna": self.skipna}
+
+
+CumulativeBlockwise._args = PatchCumulativeBlockwise._args
+CumulativeBlockwise._kwargs = PatchCumulativeBlockwise._kwargs
+
+
+def _takelast(a, skipna=True):
+    if not len(a):
+        return a
+    if skipna:
+        a = a.bfill()
+    # Cannot use `squeeze` with cudf
+    return a.tail(n=1).iloc[0]
+
+
+TakeLast.operation = staticmethod(_takelast)