Update egonet implementation

ci/cpu/build.sh

@@ -38,7 +38,7 @@ export UCX_PY_VERSION='0.29.*'
 
 # Whether to keep `dask/label/dev` channel in the env. If INSTALL_DASK_MAIN=0,
 # `dask/label/dev` channel is removed.
-export INSTALL_DASK_MAIN=0
+export INSTALL_DASK_MAIN=1
 
 ################################################################################
 # SETUP - Check environment

ci/gpu/build.sh

@@ -45,7 +45,7 @@ export UCX_PY_VERSION='0.29.*'
 
 # Whether to keep `dask/label/dev` channel in the env. If INSTALL_DASK_MAIN=0,
 # `dask/label/dev` channel is removed.
-export INSTALL_DASK_MAIN=0
+export INSTALL_DASK_MAIN=1
 
 ################################################################################
 # SETUP - Check environment

cpp/src/c_api/extract_ego.cpp

@@ -129,7 +129,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_),
-        wgt ? new cugraph::c_api::cugraph_type_erased_device_array_t(*wgt, graph_->edge_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,
                                                                data_type_id_t::SIZE_T)};

python/cugraph/cugraph/community/CMakeLists.txt

@@ -13,7 +13,7 @@
 # =============================================================================
 
 set(cython_sources
-    ecg_wrapper.pyx egonet_wrapper.pyx
+    ecg_wrapper.pyx
     ktruss_subgraph_wrapper.pyx
     leiden_wrapper.pyx
     spectral_clustering_wrapper.pyx

python/cugraph/cugraph/community/egonet.py

@@ -11,14 +11,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from cugraph.community import egonet_wrapper
-import cudf
 from cugraph.utilities import (
     ensure_cugraph_obj,
     is_nx_graph_type,
 )
 from cugraph.utilities import cugraph_to_nx
 
+import cudf
+
+from pylibcugraph import ego_graph as pylibcugraph_ego_graph
+
+from pylibcugraph import ResourceHandle
+import warnings
+
 
 def _convert_graph_to_output_type(G, input_type):
     """
@@ -44,7 +49,7 @@ def _convert_df_series_to_output_type(df, offsets, input_type):
         return df, offsets
 
 
-def ego_graph(G, n, radius=1, center=True, undirected=False, distance=None):
+def ego_graph(G, n, radius=1, center=True, undirected=None, distance=None):
     """
     Compute the induced subgraph of neighbors centered at node n,
     within a given radius.
@@ -68,10 +73,10 @@ def ego_graph(G, n, radius=1, center=True, undirected=False, distance=None):
         Defaults to True. False is not supported
 
     undirected: bool, optional
-        Defaults to False. True is not supported
+        This parameter is here for NetworkX compatibility and ignored
 
     distance: key, optional (default=None)
-        Distances are counted in hops from n. Other cases are not supported.
+        This parameter is here for NetworkX compatibility and ignored
 
     Returns
     -------
@@ -86,21 +91,54 @@ def ego_graph(G, n, radius=1, center=True, undirected=False, distance=None):
     >>> ego_graph = cugraph.ego_graph(G, 1, radius=2)
 
     """
-
     (G, input_type) = ensure_cugraph_obj(G, nx_weight_attr="weight")
-    result_graph = type(G)()
 
-    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(cudf.Series([n]))
+    result_graph = type(G)()
 
-    df, offsets = egonet_wrapper.egonet(G, n, radius)
+    if undirected is not None:
+        warning_msg = (
+            "The parameter 'undirected' is deprecated and "
+            "will be removed in the next release"
+        )
+        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 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)
+
+    # Match the seed to the vertex dtype
+    n_type = G.edgelist.edgelist_df["src"].dtype
+    n = n.astype(n_type)
+    do_expensive_check = False
+
+    source, destination, weight, _ = pylibcugraph_ego_graph(
+        resource_handle=ResourceHandle(),
+        graph=G._plc_graph,
+        source_vertices=n,
+        radius=radius,
+        do_expensive_check=do_expensive_check,
+    )
+
+    df = cudf.DataFrame()
+    df["src"] = source
+    df["dst"] = destination
+    df["weight"] = weight
 
     if G.renumbered:
         df, src_names = G.unrenumber(df, "src", get_column_names=True)
         df, dst_names = G.unrenumber(df, "dst", get_column_names=True)
+    else:
+        # FIXME: THe original 'src' and 'dst' are not stored in 'simpleGraph'
+        src_names = "src"
+        dst_names = "dst"
 
     if G.edgelist.weights:
         result_graph.from_cudf_edgelist(
@@ -111,9 +149,7 @@ def ego_graph(G, n, radius=1, center=True, undirected=False, distance=None):
     return _convert_graph_to_output_type(result_graph, input_type)
 
 
-def batched_ego_graphs(
-    G, seeds, radius=1, center=True, undirected=False, distance=None
-):
+def batched_ego_graphs(G, seeds, radius=1, center=True, undirected=None, distance=None):
     """
     Compute the induced subgraph of neighbors for each node in seeds
     within a given radius.
@@ -160,13 +196,37 @@ def batched_ego_graphs(
 
     (G, input_type) = ensure_cugraph_obj(G, nx_weight_attr="weight")
 
-    if G.renumbered is True:
-        if isinstance(seeds, cudf.DataFrame):
-            seeds = G.lookup_internal_vertex_id(seeds, seeds.columns)
-        else:
-            seeds = G.lookup_internal_vertex_id(cudf.Series(seeds))
-
-    df, offsets = egonet_wrapper.egonet(G, seeds, radius)
+    if seeds is not None:
+        if isinstance(seeds, int):
+            seeds = [seeds]
+        if isinstance(seeds, list):
+            seeds = cudf.Series(seeds)
+
+        if G.renumbered is True:
+            if isinstance(seeds, cudf.DataFrame):
+                seeds = G.lookup_internal_vertex_id(seeds, seeds.columns)
+            else:
+                seeds = G.lookup_internal_vertex_id(seeds)
+
+    # Match the seed to the vertex dtype
+    seeds_type = G.edgelist.edgelist_df["src"].dtype
+    seeds = seeds.astype(seeds_type)
+
+    do_expensive_check = False
+    source, destination, weight, offset = pylibcugraph_ego_graph(
+        resource_handle=ResourceHandle(),
+        graph=G._plc_graph,
+        source_vertices=seeds,
+        radius=radius,
+        do_expensive_check=do_expensive_check,
+    )
+
+    offsets = cudf.Series(offset)
+
+    df = cudf.DataFrame()
+    df["src"] = source
+    df["dst"] = destination
+    df["weight"] = weight
 
     if G.renumbered:
         df = G.unrenumber(df, "src", preserve_order=True)

python/cugraph/cugraph/dask/__init__.py

@@ -19,6 +19,7 @@
 from .common.read_utils import get_n_workers
 from .community.louvain import louvain
 from .community.triangle_count import triangle_count
+from .community.egonet import ego_graph
 from .centrality.katz_centrality import katz_centrality
 from .components.connectivity import weakly_connected_components
 from .sampling.uniform_neighbor_sample import uniform_neighbor_sample