Skip to content

Introduce basic "cudf" backend for Dask Expressions - #14805

Merged
rapids-bot[bot] merged 136 commits into
NVIDIA:branch-24.04from
rjzamora:new-dask-expr-backend
Mar 11, 2024
Merged

Introduce basic "cudf" backend for Dask Expressions#14805
rapids-bot[bot] merged 136 commits into
NVIDIA:branch-24.04from
rjzamora:new-dask-expr-backend

Conversation

@rjzamora

@rjzamora rjzamora commented Jan 19, 2024

Copy link
Copy Markdown
Contributor

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

@rjzamora rjzamora added 2 - In Progress Currently a work in progress improvement Improvement / enhancement to an existing function non-breaking Non-breaking change labels Jan 19, 2024
@rjzamora rjzamora self-assigned this Jan 19, 2024
@github-actions github-actions Bot added the Python Affects Python cuDF API. label Jan 19, 2024
@copy-pr-bot

copy-pr-bot Bot commented Jan 29, 2024

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@github-actions github-actions Bot added CMake CMake build issue conda labels Jan 29, 2024
@rjzamora rjzamora added 3 - Ready for Review Ready for review by team 4 - Needs Review Waiting for reviewer to review or respond and removed 2 - In Progress Currently a work in progress 3 - Ready for Review Ready for review by team labels Mar 5, 2024

@raydouglass raydouglass left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ops-codeowners changes look good

@rjzamora

Copy link
Copy Markdown
Contributor Author

/merge

@rjzamora rjzamora added 5 - Ready to Merge Testing and reviews complete, ready to merge and removed 4 - Needs Review Waiting for reviewer to review or respond labels Mar 10, 2024
@wence-

wence- commented Mar 11, 2024

Copy link
Copy Markdown
Contributor

Test fails are due to a combo of #15261 and #15265.

@bdice

bdice commented Mar 11, 2024

Copy link
Copy Markdown
Contributor

I'll merge the upstream so that the /merge command is triggered. (It was blocked by #15265)

@rjzamora

Copy link
Copy Markdown
Contributor Author

/merge

@rapids-bot
rapids-bot Bot merged commit e2fcf12 into NVIDIA:branch-24.04 Mar 11, 2024
@rjzamora
rjzamora deleted the new-dask-expr-backend branch March 11, 2024 22:04
@bdice bdice mentioned this pull request Mar 19, 2024
3 tasks
rapids-bot Bot pushed a commit that referenced this pull request Mar 20, 2024
Fixes CI blocked by dask-expr.

xref:
- #14805
- rapidsai/rapids-dask-dependency#33

Authors:
  - Bradley Dice (https://github.com/bdice)

Approvers:
  - Ray Douglass (https://github.com/raydouglass)
  - GALI PREM SAGAR (https://github.com/galipremsagar)

URL: #15343
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

5 - Ready to Merge Testing and reviews complete, ready to merge improvement Improvement / enhancement to an existing function non-breaking Non-breaking change Python Affects Python cuDF API.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants