Update egonet implementation

cpp/src/c_api/extract_ego.cpp

@@ -128,7 +128,7 @@ struct extract_ego_functor : public cugraph::c_api::abstract_functor {
 
       result_ = new cugraph::c_api::cugraph_induced_subgraph_result_t{
         new cugraph::c_api::cugraph_type_erased_device_array_t(src, graph_->vertex_type_),
-        new cugraph::c_api::cugraph_type_erased_device_array_t(dst, graph_->edge_type_),
+        new cugraph::c_api::cugraph_type_erased_device_array_t(dst, graph_->vertex_type_),
         wgt ? new cugraph::c_api::cugraph_type_erased_device_array_t(*wgt, graph_->weight_type_)
             : NULL,
         new cugraph::c_api::cugraph_type_erased_device_array_t(edge_offsets,

python/cugraph/cugraph/community/egonet.py

@@ -61,7 +61,7 @@ def ego_graph(G, n, radius=1, center=True, undirected=None, distance=None):
         information. Edge weights, if present, should be single or double
         precision floating point values.
 
-    n : integer or cudf.DataFrame
+    n : integer or list, cudf.Series, cudf.DataFrame
         A single node as integer or a cudf.DataFrame if nodes are
         represented with multiple columns. If a cudf.DataFrame is provided,
         only the first row is taken as the node input.
@@ -73,10 +73,10 @@ def ego_graph(G, n, radius=1, center=True, undirected=None, distance=None):
         Defaults to True. False is not supported
 
     undirected: bool, optional
-        This parameter is here for NetworkX compatibility and ignored
+        This parameter is here for NetworkX compatibility and is ignored
 
     distance: key, optional (default=None)
-        This parameter is here for NetworkX compatibility and ignored
+        This parameter is here for NetworkX compatibility and is ignored
 
     Returns
     -------
@@ -93,7 +93,7 @@ def ego_graph(G, n, radius=1, center=True, undirected=None, distance=None):
     """
     (G, input_type) = ensure_cugraph_obj(G, nx_weight_attr="weight")
 
-    result_graph = type(G)()
+    result_graph = type(G)(directed=G.is_directed())
 
     if undirected is not None:
         warning_msg = (
@@ -102,17 +102,21 @@ def ego_graph(G, n, radius=1, center=True, undirected=None, distance=None):
         )
         warnings.warn(warning_msg, PendingDeprecationWarning)
 
-    if n is not None:
-        if isinstance(n, int):
-            n = [n]
-        if isinstance(n, list):
-            n = cudf.Series(n)
-
+    if isinstance(n, int):
+        n = [n]
+    if isinstance(n, list):
+        n = cudf.Series(n)
+    if isinstance(n, cudf.Series):
         if G.renumbered is True:
-            if isinstance(n, cudf.DataFrame):
-                n = G.lookup_internal_vertex_id(n, n.columns)
-            else:
-                n = G.lookup_internal_vertex_id(n)
+            n = G.lookup_internal_vertex_id(n)
+    elif isinstance(n, cudf.DataFrame):
+        if G.renumbered is True:
+            n = G.lookup_internal_vertex_id(n, n.columns)
+    else:
+        raise TypeError(
+            f"'n' must be either an integer or a list or a cudf.Series"
+            f" or a cudf.DataFrame, got: {type(n)}"
+        )
 
     # Match the seed to the vertex dtype
     n_type = G.edgelist.edgelist_df["src"].dtype

python/cugraph/cugraph/dask/community/egonet.py

@@ -41,6 +41,17 @@ def _call_ego_graph(
 
 
 def consolidate_results(ddf, offsets, num_seeds):
+    """
+    Each rank returns its ego_graph dataframe with its corresponding
+    offsets array. This is ideal if the user operates on distributed memory
+    but when attempting to bring the result into a single machine,
+    the ego_graph dataframes generated from each seed cannot be extracted
+    using the offsets array. This function consolidate the final result by
+    performing segmented copies.
+
+    Returns: consolidated ego_graph dataframe and offsets array
+    """
+
     df = cudf.DataFrame()
     offset_array = [0]
     for s in range(num_seeds):
@@ -94,7 +105,7 @@ def ego_graph(input_graph, n, radius=1, center=True):
         information. Edge weights, if present, should be single or double
         precision floating point values.
 
-    n : int, list or cudf.Series, cudf.DataFrame
+    n : int, list or cudf Series or Dataframe, dask_cudf Series or DataFrame
         A node or a list or cudf.Series of nodes or a cudf.DataFrame if nodes
         are represented with multiple columns. If a cudf.DataFrame is provided,
         only the first row is taken as the node input.
@@ -119,25 +130,29 @@ def ego_graph(input_graph, n, radius=1, center=True):
     # Initialize dask client
     client = input_graph._client
 
-    if n is not None:
-        if isinstance(n, int):
-            n = [n]
-        if isinstance(n, list):
-            n = cudf.Series(n)
-        if not isinstance(n, cudf.Series):
+    if isinstance(n, int):
+        n = [n]
+    if isinstance(n, list):
+        n = cudf.Series(n)
+    elif not isinstance(n, (cudf.Series, dask_cudf.Series)):
+        if not isinstance(n, (cudf.DataFrame, dask_cudf.DataFrame)):
             raise TypeError(
-                f"'n' must be either a list or a cudf.Series," f"got: {n.dtype}"
+                f"'n' must be either an integer or a list or a "
+                f"cudf or dask_cudf Series or DataFrame, got: {type(n)}"
             )
-        num_seeds = len(n)
-        # n uses "external" vertex IDs, but since the graph has been
-        # renumbered, the node ID must also be renumbered.
-        if input_graph.renumbered:
-            n = input_graph.lookup_internal_vertex_id(n).compute()
-            n_type = input_graph.edgelist.edgelist_df.dtypes[0]
-        else:
-            n_type = input_graph.input_df.dtypes[0]
-
-    n = dask_cudf.from_cudf(n, npartitions=min(input_graph._npartitions, len(n)))
+
+    num_seeds = len(n)
+    # n uses "external" vertex IDs, but since the graph has been
+    # renumbered, the node ID must also be renumbered.
+    if input_graph.renumbered:
+        n = input_graph.lookup_internal_vertex_id(n)
+        n_type = input_graph.edgelist.edgelist_df.dtypes[0]
+    else:
+        n_type = input_graph.input_df.dtypes[0]
+
+    if isinstance(n, (cudf.Series, cudf.DataFrame)):
+        n = dask_cudf.from_cudf(n, npartitions=min(input_graph._npartitions, len(n)))
+
     n = n.astype(n_type)
 
     n = get_distributed_data(n)

python/cugraph/cugraph/tests/mg/test_mg_egonet.py

@@ -16,6 +16,7 @@
 import pytest
 import cugraph
 import dask_cudf
+from cugraph.dask.common.mg_utils import is_single_gpu
 
 # from cugraph.dask.common.mg_utils import is_single_gpu
 from cugraph.testing import utils
@@ -113,9 +114,7 @@ def input_expected_output(input_combo):
 # =============================================================================
 
 
-# @pytest.mark.skipif(
-#    is_single_gpu(), reason="skipping MG testing on Single GPU system"
-# )
+@pytest.mark.skipif(is_single_gpu(), reason="skipping MG testing on Single GPU system")
 def test_dask_ego_graphs(dask_client, benchmark, input_expected_output):
 
     dg = input_expected_output["MGGraph"]

python/pylibcugraph/pylibcugraph/egonet.pyx

@@ -83,7 +83,6 @@ def ego_graph(ResourceHandle resource_handle,
     radius: size_t
         The number of hops to go out from each source vertex
 
-
     do_expensive_check : bool_t
         If True, performs more extensive tests on the inputs to ensure
         validitity, at the expense of increased run time.
@@ -153,15 +152,14 @@ def ego_graph(ResourceHandle resource_handle,
     cdef cugraph_type_erased_device_array_view_t* subgraph_offsets_ptr = \
         cugraph_induced_subgraph_get_subgraph_offsets(result_ptr)
 
+    # FIXME: Get ownership of the result data instead of performing a copy
+    # for perfomance improvement
     cupy_sources = copy_to_cupy_array(
         c_resource_handle_ptr, sources_ptr)
-
     cupy_destinations = copy_to_cupy_array(
         c_resource_handle_ptr, destinations_ptr)
-
     cupy_edge_weights = copy_to_cupy_array(
         c_resource_handle_ptr, edge_weights_ptr)
-
     cupy_subgraph_offsets = copy_to_cupy_array(
         c_resource_handle_ptr, subgraph_offsets_ptr)