diff --git a/python/cugraph/cugraph/structure/convert_matrix.py b/python/cugraph/cugraph/structure/convert_matrix.py index d905866689c..62a97113816 100644 --- a/python/cugraph/cugraph/structure/convert_matrix.py +++ b/python/cugraph/cugraph/structure/convert_matrix.py @@ -309,7 +309,7 @@ def from_pandas_edgelist(df, return G -def to_pandas_edgelist(G, source='source', destination='destination'): +def to_pandas_edgelist(G, source='src', destination='dst'): """ Returns the graph edge list as a Pandas DataFrame. diff --git a/python/cugraph/cugraph/structure/graph_implementation/simpleGraph.py b/python/cugraph/cugraph/structure/graph_implementation/simpleGraph.py index 2b57736562b..3475e8cbbd1 100644 --- a/python/cugraph/cugraph/structure/graph_implementation/simpleGraph.py +++ b/python/cugraph/cugraph/structure/graph_implementation/simpleGraph.py @@ -182,16 +182,19 @@ def __from_edgelist( if self.batch_enabled: self._replicate_edgelist() - def to_pandas_edgelist(self, source='source', destination='destination'): + def to_pandas_edgelist(self, source='src', destination='dst', + weight='weights'): """ Returns the graph edge list as a Pandas DataFrame. Parameters ---------- - source : str or array-like, optional (default='source') + source : str or array-like, optional (default='src') source column name or array of column names - destination : str or array-like, optional (default='destination') + destination : str or array-like, optional (default='dst') destination column name or array of column names + weight : str or array-like, optional (default='weight') + weight column name or array of column names Returns ------- @@ -199,6 +202,12 @@ def to_pandas_edgelist(self, source='source', destination='destination'): """ gdf = self.view_edge_list() + if self.properties.weighted: + gdf.rename(columns={'src': source, 'dst': destination, + 'weight': weight}, inplace=True) + else: + gdf.rename(columns={'src': source, + 'dst': destination}, inplace=True) return gdf.to_pandas() def to_pandas_adjacency(self): diff --git a/python/cugraph/cugraph/tests/test_convert_matrix.py b/python/cugraph/cugraph/tests/test_convert_matrix.py index 9d73b756f61..0cd9061883d 100644 --- a/python/cugraph/cugraph/tests/test_convert_matrix.py +++ b/python/cugraph/cugraph/tests/test_convert_matrix.py @@ -53,7 +53,7 @@ def test_to_from_pandas(graph_file): # create a cugraph DiGraph and convert to pandas adjacency cuG = cugraph.from_pandas_edgelist( M, source="0", destination="1", edge_attr="weight", - create_using=cugraph.DiGraph + create_using=cugraph.Graph(directed=True) ) cu_pdf = cugraph.to_pandas_adjacency(cuG) @@ -65,8 +65,9 @@ def test_to_from_pandas(graph_file): # Convert pandas adjacency list to graph new_nxG = nx.from_pandas_adjacency(nx_pdf, create_using=nx.DiGraph) - new_cuG = cugraph.from_pandas_adjacency(cu_pdf, - create_using=cugraph.DiGraph) + new_cuG = cugraph.from_pandas_adjacency( + cu_pdf, + create_using=cugraph.Graph(directed=True)) # Compare pandas edgelist exp_pdf = nx.to_pandas_edgelist(new_nxG) diff --git a/python/cugraph/cugraph/tests/test_graph.py b/python/cugraph/cugraph/tests/test_graph.py index e2e619cfa55..b6c285fb0a7 100644 --- a/python/cugraph/cugraph/tests/test_graph.py +++ b/python/cugraph/cugraph/tests/test_graph.py @@ -666,6 +666,18 @@ def test_neighbors(graph_file): assert cu_neighbors == nx_neighbors +# Test +@pytest.mark.parametrize("graph_file", utils.DATASETS) +def test_to_pandas_edgelist(graph_file): + cu_M = utils.read_csv_file(graph_file) + + G = cugraph.Graph() + G.from_cudf_edgelist(cu_M, source="0", destination="1") + + assert 's' in G.to_pandas_edgelist('s', 'd').columns + assert 's' in G.to_pandas_edgelist(source='s', destination='d').columns + + def test_graph_init_with_multigraph(): """ Ensures only a valid MultiGraph instance can be used to initialize a Graph