Migrate SG and MG BFS to pylibcugraph (#2284)

Closes #2270

Authors:
  - Alex Barghi (https://github.com/alexbarghi-nv)
  - Chuck Hastings (https://github.com/ChuckHastings)

Approvers:
  - Rick Ratzel (https://github.com/rlratzel)
  - https://github.com/betochimas

URL: #2284

.gitignore

@@ -88,3 +88,4 @@ python/cugraph/cugraph/tests/dask-worker-space
 
 # Sphinx docs & build artifacts
 docs/cugraph/source/api_docs/api/*
+

python/cugraph/cugraph/dask/traversal/bfs.py

@@ -13,45 +13,75 @@
 # limitations under the License.
 #
 
-from collections.abc import Iterable
+from pylibcugraph.experimental import (MGGraph,
+                                       ResourceHandle,
+                                       GraphProperties,
+                                       )
+from pylibcugraph import bfs as pylibcugraph_bfs
 
 from dask.distributed import wait, default_client
-from cugraph.dask.common.input_utils import (get_distributed_data,
-                                             get_vertex_partition_offsets)
-from cugraph.dask.traversal import mg_bfs_wrapper as mg_bfs
+from cugraph.dask.common.input_utils import get_distributed_data
 import cugraph.dask.comms.comms as Comms
 import cudf
 import dask_cudf
 
 
-def call_bfs(sID,
-             data,
-             src_col_name,
-             dst_col_name,
-             num_verts,
-             num_edges,
-             vertex_partition_offsets,
-             aggregate_segment_offsets,
-             start,
-             depth_limit,
-             return_distances):
-    wid = Comms.get_worker_id(sID)
-    handle = Comms.get_handle(sID)
-    local_size = len(aggregate_segment_offsets) // Comms.get_n_workers(sID)
-    segment_offsets = \
-        aggregate_segment_offsets[local_size * wid: local_size * (wid + 1)]
-    return mg_bfs.mg_bfs(data[0],
-                         src_col_name,
-                         dst_col_name,
-                         num_verts,
-                         num_edges,
-                         vertex_partition_offsets,
-                         wid,
-                         handle,
-                         segment_offsets,
-                         start,
-                         depth_limit,
-                         return_distances)
+def _call_plc_mg_bfs(
+                    sID,
+                    data,
+                    sources,
+                    depth_limit,
+                    src_col_name,
+                    dst_col_name,
+                    graph_properties,
+                    num_edges,
+                    direction_optimizing=False,
+                    do_expensive_check=False,
+                    return_predecessors=True):
+    comms_handle = Comms.get_handle(sID)
+    resource_handle = ResourceHandle(comms_handle.getHandle())
+
+    srcs = cudf.Series(data[0][src_col_name], dtype='int32')
+    dsts = cudf.Series(data[0][dst_col_name], dtype='int32')
+    weights = cudf.Series(data[0]['value'], dtype='float32') \
+        if 'value' in data[0].columns \
+        else cudf.Series((srcs + 1) / (srcs + 1), dtype='float32')
+
+    mg = MGGraph(
+        resource_handle=resource_handle,
+        graph_properties=graph_properties,
+        src_array=srcs,
+        dst_array=dsts,
+        weight_array=weights,
+        store_transposed=False,
+        num_edges=num_edges,
+        do_expensive_check=do_expensive_check
+    )
+
+    res = \
+        pylibcugraph_bfs(
+            resource_handle,
+            mg,
+            cudf.Series(sources, dtype='int32'),
+            direction_optimizing,
+            depth_limit if depth_limit is not None else 0,
+            return_predecessors,
+            True
+        )
+
+    return res
+
+
+def convert_to_cudf(cp_arrays):
+    """
+    create a cudf DataFrame from cupy arrays
+    """
+    cupy_distances, cupy_predecessors, cupy_vertices = cp_arrays
+    df = cudf.DataFrame()
+    df["vertex"] = cupy_vertices
+    df["distance"] = cupy_distances
+    df["predecessor"] = cupy_predecessors
+    return df
 
 
 def bfs(input_graph,
@@ -115,95 +145,81 @@ def bfs(input_graph,
 
     input_graph.compute_renumber_edge_list(transposed=False)
     ddf = input_graph.edgelist.edgelist_df
-    vertex_partition_offsets = get_vertex_partition_offsets(input_graph)
-    num_verts = vertex_partition_offsets.iloc[-1]
+
+    graph_properties = GraphProperties(
+        is_multigraph=False)
 
     num_edges = len(ddf)
     data = get_distributed_data(ddf)
 
+    src_col_name = input_graph.renumber_map.renumbered_src_col_name
+    dst_col_name = input_graph.renumber_map.renumbered_dst_col_name
+    renumber_ddf = input_graph.renumber_map.implementation.ddf
+    col_names = input_graph.renumber_map.implementation.col_names
+
+    if isinstance(start, dask_cudf.DataFrame) \
+            or isinstance(start, cudf.DataFrame):
+        tmp_df = start
+        tmp_col_names = start.columns
+    else:
+        tmp_df = cudf.DataFrame()
+        tmp_df["0"] = cudf.Series(start)
+        tmp_col_names = ["0"]
+
+    original_start_len = len(tmp_df)
+
+    tmp_ddf = tmp_df[tmp_col_names].rename(
+        columns=dict(zip(tmp_col_names, col_names)))
+    for name in col_names:
+        tmp_ddf[name] = tmp_ddf[name].astype(renumber_ddf[name].dtype)
+    renumber_data = get_distributed_data(renumber_ddf)
+
     def df_merge(df, tmp_df, tmp_col_names):
         x = df[0].merge(tmp_df, on=tmp_col_names, how='inner')
         return x['global_id']
 
-    if input_graph.renumbered:
-        src_col_name = input_graph.renumber_map.renumbered_src_col_name
-        dst_col_name = input_graph.renumber_map.renumbered_dst_col_name
-        renumber_ddf = input_graph.renumber_map.implementation.ddf
-        col_names = input_graph.renumber_map.implementation.col_names
-        if isinstance(start,
-                      dask_cudf.DataFrame) or isinstance(start,
-                                                         cudf.DataFrame):
-            tmp_df = start
-            tmp_col_names = start.columns
-        else:
-            tmp_df = cudf.DataFrame()
-            tmp_df["0"] = cudf.Series(start)
-            tmp_col_names = ["0"]
-
-        original_start_len = len(tmp_df)
-
-        tmp_ddf = tmp_df[tmp_col_names].rename(
-            columns=dict(zip(tmp_col_names, col_names)))
-        for name in col_names:
-            tmp_ddf[name] = tmp_ddf[name].astype(renumber_ddf[name].dtype)
-        renumber_data = get_distributed_data(renumber_ddf)
-        start = [client.submit(df_merge,
-                               wf[1],
-                               tmp_ddf,
-                               col_names,
-                               workers=[wf[0]])
-                 for idx, wf in enumerate(renumber_data.worker_to_parts.items()
-                                          )
-                 ]
-
-        def count_src(df):
-            return len(df)
-
-        count_src_results = client.map(count_src, start)
-        cg = client.gather(count_src_results)
-        if sum(cg) < original_start_len:
-            raise ValueError('At least one start vertex provided was invalid')
-
-    else:
-        # If the input graph was created with renumbering disabled (Graph(...,
-        # renumber=False), the above compute_renumber_edge_list() call will not
-        # perform a renumber step and the renumber_map will not have src/dst
-        # col names. In that case, the src/dst values specified when reading
-        # the edgelist dataframe are to be used, but only if they were single
-        # string values (ie. not a list representing multi-columns).
-        if isinstance(input_graph.source_columns, Iterable):
-            raise RuntimeError("input_graph was not renumbered but has a "
-                               "non-string source column name (got: "
-                               f"{input_graph.source_columns}). Re-create "
-                               "input_graph with either renumbering enabled "
-                               "or a source column specified as a string.")
-        if isinstance(input_graph.destination_columns, Iterable):
-            raise RuntimeError("input_graph was not renumbered but has a "
-                               "non-string destination column name (got: "
-                               f"{input_graph.destination_columns}). "
-                               "Re-create input_graph with either renumbering "
-                               "enabled or a destination column specified as "
-                               "a string.")
-        src_col_name = input_graph.source_columns
-        dst_col_name = input_graph.destination_columns
-
-    result = [client.submit(
-              call_bfs,
+    start = [
+        client.submit(
+            df_merge,
+            wf[1],
+            tmp_ddf,
+            col_names,
+            workers=[wf[0]]
+        )
+        for idx, wf in enumerate(renumber_data.worker_to_parts.items())
+    ]
+
+    def count_src(df):
+        return len(df)
+
+    count_src_results = client.map(count_src, start)
+    cg = client.gather(count_src_results)
+    if sum(cg) < original_start_len:
+        raise ValueError('At least one start vertex provided was invalid')
+
+    cupy_result = [client.submit(
+              _call_plc_mg_bfs,
               Comms.get_session_id(),
               wf[1],
+              start[idx],
+              depth_limit,
               src_col_name,
               dst_col_name,
-              num_verts,
+              graph_properties,
               num_edges,
-              vertex_partition_offsets,
-              input_graph.aggregate_segment_offsets,
-              start[idx],
-              depth_limit,
+              False,
+              True,
               return_distances,
               workers=[wf[0]])
               for idx, wf in enumerate(data.worker_to_parts.items())]
-    wait(result)
-    ddf = dask_cudf.from_delayed(result)
+    wait(cupy_result)
+
+    cudf_result = [client.submit(convert_to_cudf,
+                                 cp_arrays)
+                   for cp_arrays in cupy_result]
+    wait(cudf_result)
+
+    ddf = dask_cudf.from_delayed(cudf_result)
 
     if input_graph.renumbered:
         ddf = input_graph.unrenumber(ddf, 'vertex')