Enable expression-based Dask Dataframe support (#4325)

**[WIP]** I'm using this PR to debug/add support for `DASK_DATAFRAME__QUERY_PLANNING=True`.

**NOTES**:
- Depends on dask/dask-expr#1041 [Merged]
- Depends on dask/dask-expr#1044

Authors:
  - Richard (Rick) Zamora (https://github.com/rjzamora)

Approvers:
  - Rick Ratzel (https://github.com/rlratzel)
  - Ray Douglass (https://github.com/raydouglass)

URL: #4325

Author: rjzamora

benchmarks/cugraph/standalone/bulk_sampling/cugraph_bulk_sampling.py

@@ -344,7 +344,7 @@ def generate_rmat_dataset(
     del label_df
     gc.collect()
 
-    dask_label_df = dask_cudf.from_dask_dataframe(dask_label_df)
+    dask_label_df = dask_label_df.to_backend("cudf")
 
     node_offsets = {"paper": 0}
     edge_offsets = {("paper", "cites", "paper"): 0}

ci/test_python.sh

@@ -3,10 +3,6 @@
 
 set -euo pipefail
 
-# TODO: Enable dask query planning (by default) once some bugs are fixed.
-# xref: https://github.com/rapidsai/cudf/issues/15027
-export DASK_DATAFRAME__QUERY_PLANNING=False
-
 # Support invoking test_python.sh outside the script directory
 cd "$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/../
 

ci/test_wheel.sh

@@ -3,10 +3,6 @@
 
 set -eoxu pipefail
 
-# TODO: Enable dask query planning (by default) once some bugs are fixed.
-# xref: https://github.com/rapidsai/cudf/issues/15027
-export DASK_DATAFRAME__QUERY_PLANNING=False
-
 package_name=$1
 package_dir=$2
 

python/cugraph/cugraph/dask/__init__.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2023, NVIDIA CORPORATION.
+# Copyright (c) 2020-2024, NVIDIA CORPORATION.
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
@@ -11,6 +11,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from dask import config
+
 from .link_analysis.pagerank import pagerank
 from .link_analysis.hits import hits
 from .traversal.bfs import bfs
@@ -34,3 +36,6 @@
 from .link_prediction.sorensen import sorensen
 from .link_prediction.overlap import overlap
 from .community.leiden import leiden
+
+# Avoid "p2p" shuffling in dask for now
+config.set({"dataframe.shuffle.method": "tasks"})

python/cugraph/cugraph/dask/common/input_utils.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2023, NVIDIA CORPORATION.
+# Copyright (c) 2020-2024, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,8 +15,8 @@
 
 from collections.abc import Sequence
 from collections import OrderedDict
-from dask_cudf.core import DataFrame as dcDataFrame
-from dask_cudf.core import Series as daskSeries
+from dask_cudf import DataFrame as dcDataFrame
+from dask_cudf import Series as daskSeries
 
 import cugraph.dask.comms.comms as Comms
 

python/cugraph/cugraph/dask/common/part_utils.py

@@ -18,8 +18,8 @@
 import collections
 import dask_cudf
 from dask.array.core import Array as daskArray
-from dask_cudf.core import DataFrame as daskDataFrame
-from dask_cudf.core import Series as daskSeries
+from dask_cudf import DataFrame as daskDataFrame
+from dask_cudf import Series as daskSeries
 from functools import reduce
 import cugraph.dask.comms.comms as Comms
 from dask.delayed import delayed

python/cugraph/cugraph/structure/convert_matrix.py

@@ -40,7 +40,7 @@ def from_edgelist(
 
     Parameters
     ----------
-    df : cudf.DataFrame, pandas.DataFrame, dask_cudf.core.DataFrame
+    df : cudf.DataFrame, pandas.DataFrame, dask_cudf.DataFrame
         This DataFrame contains columns storing edge source vertices,
         destination (or target following NetworkX's terminology) vertices, and
         (optional) weights.
@@ -95,7 +95,7 @@ def from_edgelist(
             renumber=renumber,
         )
 
-    elif df_type is dask_cudf.core.DataFrame:
+    elif df_type is dask_cudf.DataFrame:
         if create_using is None:
             G = Graph()
         elif isinstance(create_using, Graph):

python/cugraph/cugraph/structure/graph_implementation/simpleDistributedGraph.py

@@ -285,19 +285,20 @@ def __from_edgelist(
                 symmetrize=not self.properties.directed,
             )
 
+        # Create a dask_cudf dataframe from the cudf series
+        # or dataframe objects obtained from symmetrization
         if isinstance(source_col, dask_cudf.Series):
-            # Create a dask_cudf dataframe from the cudf series obtained
-            # from symmetrization
-            input_ddf = source_col.to_frame()
-            input_ddf = input_ddf.rename(columns={source_col.name: source})
-            input_ddf[destination] = dest_col
+            frames = [
+                source_col.to_frame(name=source),
+                dest_col.to_frame(name=destination),
+            ]
         else:
-            # Multi column dask_cudf dataframe
-            input_ddf = dask_cudf.concat([source_col, dest_col], axis=1)
+            frames = [source_col, dest_col]
 
         if value_col is not None:
-            for vc in value_col_names:
-                input_ddf[vc] = value_col[vc]
+            frames.append(value_col[value_col_names])
+
+        input_ddf = dask_cudf.concat(frames, axis=1)
 
         self.input_df = input_ddf
 

python/cugraph/cugraph/tests/data_store/test_property_graph_mg.py

@@ -159,8 +159,8 @@ def df_type_id(dataframe_type):
         return s + "cudf.DataFrame"
     if dataframe_type == pd.DataFrame:
         return s + "pandas.DataFrame"
-    if dataframe_type == dask_cudf.core.DataFrame:
-        return s + "dask_cudf.core.DataFrame"
+    if dataframe_type == dask_cudf.DataFrame:
+        return s + "dask_cudf.DataFrame"
     return s + "?"
 
 

python/cugraph/cugraph/tests/internals/test_symmetrize_mg.py

@@ -232,14 +232,17 @@ def test_mg_symmetrize(dask_client, read_datasets):
 
     # create a dask DataFrame from the dask Series
     if isinstance(sym_src, dask_cudf.Series):
-        ddf2 = sym_src.to_frame()
-        ddf2 = ddf2.rename(columns={sym_src.name: "src"})
-        ddf2["dst"] = sym_dst
+        frames = [
+            sym_src.to_frame(name="src"),
+            sym_dst.to_frame(name="dst"),
+        ]
     else:
-        ddf2 = dask_cudf.concat([sym_src, sym_dst], axis=1)
+        frames = [sym_src, sym_dst]
 
     if val_col_name is not None:
-        ddf2["weight"] = sym_val
+        frames.append(sym_val.to_frame(name="weight"))
+
+    ddf2 = dask_cudf.concat(frames, axis=1)
 
     compare(ddf, ddf2, src_col_name, dst_col_name, val_col_name)
 

python/cugraph/cugraph/tests/structure/test_graph_mg.py

@@ -99,13 +99,13 @@ def test_nodes_functionality(dask_client, input_combo):
     expected_nodes = (
         dask_cudf.concat([ddf["src"], ddf["dst"]])
         .drop_duplicates()
-        .to_frame()
-        .sort_values(0)
+        .to_frame(name="0")
+        .sort_values("0")
     )
 
     expected_nodes = expected_nodes.compute().reset_index(drop=True)
 
-    result_nodes["expected_nodes"] = expected_nodes[0]
+    result_nodes["expected_nodes"] = expected_nodes["0"]
 
     compare = result_nodes.query("result_nodes != expected_nodes")
 

python/cugraph/cugraph/tests/utils/test_dataset.py

@@ -198,7 +198,7 @@ def test_reader_dask(dask_client, dataset):
     E = dataset.get_dask_edgelist(download=True)
 
     assert E is not None
-    assert isinstance(E, dask_cudf.core.DataFrame)
+    assert isinstance(E, dask_cudf.DataFrame)
     dataset.unload()