add batched ego test

rapidsai · Feb 3, 2021 · 3a3aa89 · 3a3aa89
1 parent ae4b7cb
commit 3a3aa89
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 4 deletions.
diff --git a/python/cugraph/__init__.py b/python/cugraph/__init__.py
@@ -25,6 +25,7 @@
     subgraph,
     triangles,
     ego_graph,
+    batched_ego_graphs,
 )
 
 from cugraph.structure import (

diff --git a/python/cugraph/community/__init__.py b/python/cugraph/community/__init__.py
@@ -26,3 +26,4 @@
 from cugraph.community.ktruss_subgraph import ktruss_subgraph
 from cugraph.community.ktruss_subgraph import k_truss
 from cugraph.community.egonet import ego_graph
+from cugraph.community.egonet import batched_ego_graphs
diff --git a/python/cugraph/community/egonet.py b/python/cugraph/community/egonet.py
@@ -149,7 +149,7 @@ def batched_ego_graphs(
     df, offsets = egonet_wrapper.egonet(G, seeds, radius)
 
     if G.renumbered:
-        df = G.unrenumber(df, "src")
-        df = G.unrenumber(df, "dst")
+        df = G.unrenumber(df, "src", preserve_order=True)
+        df = G.unrenumber(df, "dst", preserve_order=True)
 
     return _convert_df_series_to_output_type(df, offsets, input_type)
diff --git a/python/cugraph/structure/number_map.py b/python/cugraph/structure/number_map.py
@@ -895,7 +895,7 @@ def unrenumber(self, df, column_name, preserve_order=False):
         if preserve_order:
             df = df.sort_values(
                 index_name
-            ).drop(index_name).reset_index(drop=True)
+            ).drop(columns=index_name).reset_index(drop=True)
 
         if type(df) is dask_cudf.DataFrame:
             return df.map_partitions(

diff --git a/python/cugraph/tests/test_egonet.py b/python/cugraph/tests/test_egonet.py
@@ -54,7 +54,10 @@ def test_ego_graph_nx(graph_file, seed, radius):
     assert nx.is_isomorphic(ego_nx, ego_cugraph)
 
 
-def batched_ego_graphs(G, seeds, radius=1):
+@pytest.mark.parametrize("graph_file", utils.DATASETS)
+@pytest.mark.parametrize("seeds", [[0, 5, 13]])
+@pytest.mark.parametrize("radius", [1, 2, 3])
+def test_batched_ego_graphs(graph_file, seeds, radius):
     """
     Compute the  induced subgraph of neighbors for each node in seeds
     within a given radius.
@@ -78,3 +81,21 @@ def batched_ego_graphs(G, seeds, radius=1):
         Series containing the starting offset in the returned edge list
         for each seed.
     """
+    gc.collect()
+
+    # Nx
+    df = utils.read_csv_for_nx(graph_file, read_weights_in_sp=True)
+    Gnx = nx.from_pandas_edgelist(
+        df, create_using=nx.Graph(), source="0", target="1", edge_attr="weight"
+    )
+
+    # cugraph
+    df, offsets = cugraph.batched_ego_graphs(Gnx, seeds, radius=radius)
+    for i in range(len(seeds)):
+        ego_nx = nx.ego_graph(Gnx, seeds[i], radius=radius)
+        ego_df = df[offsets[i]:offsets[i+1]]
+        ego_cugraph = nx.from_pandas_edgelist(ego_df,
+                                              source="src",
+                                              target="dst",
+                                              edge_attr="weight")
+    assert nx.is_isomorphic(ego_nx, ego_cugraph)