Description

Mostly addresses #15027

dask/dask-expr#728 exposed the necessary mechanisms for us to define a custom dask-expr backend for cudf. The new dispatching mechanisms are effectively the same as those in dask.dataframe. The only difference is that we are now registering/implementing "expression-based" collections.

This PR does the following:

• Defines a basic DataFrameBackendEntrypoint class for collection creation, and registers new collections using get_collection_type.
• Refactors the dask_cudf import structure to properly support the "dataframe.query-planning" configuration.
• Modifies CI to test dask-expr support for some of the dask_cudf tests. This coverage can be expanded in follow-up work.

Experimental Change: This PR patches dask_expr._expr.Expr.__new__ to enable type-based dispatching. This effectively allows us to surgically replace problematic Expr subclasses that do not work for cudf-backed data. For example, this PR replaces the upstream TakeLast expression to avoid using squeeze (since this method is not supported by cudf). This particular fix can be moved upstream relatively easily. However, having this kind of "patching" mechanism may be valuable for more complicated pandas/cudf discrepancies.

Usage example

from dask import config
config.set({"dataframe.query-planning": True})
import dask_cudf

df = dask_cudf.DataFrame.from_dict(
    {"x": range(100), "y":  [1, 2, 3, 4] * 25, "z": ["1", "2"] * 50},
    npartitions=10,
)
df["y2"] = df["x"] + df["y"]
agg = df.groupby("y").agg({"y2": "mean"})["y2"]
agg.simplify().pprint()

Dask cuDF should now be using dask-expr for "query planning":

Projection: columns='y2'
  GroupbyAggregation: arg={'y2': 'mean'} observed=True split_out=1'y'
    Assign: y2=
      Projection: columns=['y']
        FromPandas: frame='<dataframe>' npartitions=10 columns=['x', 'y']
      Add:
        Projection: columns='x'
          FromPandas: frame='<dataframe>' npartitions=10 columns=['x', 'y']
        Projection: columns='y'
          FromPandas: frame='<dataframe>' npartitions=10 columns=['x', 'y']

TODO

• Add basic tests
• Confirm that general design makes sense

Follow Up Work:

• Expand dask-expr test coverage
• Fix local and upstream bugs
• Add documentation once "critical mass" is reached