Skip to content

Commit

Permalink
Adding metadata getter methods to datasets API (#3821)
Browse files Browse the repository at this point in the history
Closes #3820 

This PR adds simple getter methods to the `dataset` class, which allows users to easily get information about datasets without need to access the `metadata` dict or look in the directory.

```python
from cugraph.datasets import karate

# users now call
karate.number_of_nodes()

# instead of 
karate.metadata['number_of_nodes']
```

Authors:
  - ralph (https://github.com/nv-rliu)

Approvers:
  - Alex Barghi (https://github.com/alexbarghi-nv)

URL: #3821
  • Loading branch information
nv-rliu authored Sep 8, 2023
1 parent 4ee227c commit 6779e89
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 0 deletions.
36 changes: 36 additions & 0 deletions python/cugraph/cugraph/datasets/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -266,6 +266,42 @@ def get_path(self):

return self._path.absolute()

def is_directed(self):
"""
Returns True if the graph is a directed graph.
"""
return self.metadata["is_directed"]

def is_multigraph(self):
"""
Returns True if the graph is a multigraph.
"""
return self.metadata["is_multigraph"]

def is_symmetric(self):
"""
Returns True if the graph is symmetric.
"""
return self.metadata["is_symmetric"]

def number_of_nodes(self):
"""
An alias of number_of_vertices()
"""
return self.number_of_vertices()

def number_of_vertices(self):
"""
Get the number of vertices in the graph.
"""
return self.metadata["number_of_nodes"]

def number_of_edges(self):
"""
Get the number of edges in the graph.
"""
return self.metadata["number_of_edges"]


def download_all(force=False):
"""
Expand Down
10 changes: 10 additions & 0 deletions python/cugraph/cugraph/tests/utils/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -328,6 +328,16 @@ def test_is_multigraph(dataset):
assert G.is_multigraph() == dataset.metadata["is_multigraph"]


@pytest.mark.parametrize("dataset", ALL_DATASETS)
def test_object_getters(dataset):
assert dataset.is_directed() == dataset.metadata["is_directed"]
assert dataset.is_multigraph() == dataset.metadata["is_multigraph"]
assert dataset.is_symmetric() == dataset.metadata["is_symmetric"]
assert dataset.number_of_nodes() == dataset.metadata["number_of_nodes"]
assert dataset.number_of_vertices() == dataset.metadata["number_of_nodes"]
assert dataset.number_of_edges() == dataset.metadata["number_of_edges"]


#
# Test experimental for DeprecationWarnings
#
Expand Down

0 comments on commit 6779e89

Please sign in to comment.