update test

rapidsai · Nov 7, 2023 · f619931 · f619931
1 parent d7eb7d7
commit f619931
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 5 deletions.
diff --git a/python/cugraph-pyg/cugraph_pyg/data/cugraph_store.py b/python/cugraph-pyg/cugraph_pyg/data/cugraph_store.py
@@ -507,8 +507,6 @@ def __construct_graph(
             if df.etp.dtype != "int32":
                 raise ValueError("Edge type must be int32!")
 
-            print(df)
-
             # Ensure the dataframe is constructed on each partition
             # instead of adding additional synchronization head from potential
             # host to device copies.

diff --git a/python/cugraph-pyg/cugraph_pyg/tests/mg/test_mg_cugraph_loader.py b/python/cugraph-pyg/cugraph_pyg/tests/mg/test_mg_cugraph_loader.py
@@ -15,7 +15,6 @@
 
 from cugraph_pyg.loader import CuGraphNeighborLoader
 from cugraph_pyg.data import CuGraphStore
-
 from cugraph.utilities.utils import import_optional, MissingModule
 
 torch = import_optional("torch")

diff --git a/python/cugraph-pyg/cugraph_pyg/tests/mg/test_mg_cugraph_store.py b/python/cugraph-pyg/cugraph_pyg/tests/mg/test_mg_cugraph_store.py
@@ -120,7 +120,7 @@ def test_get_edge_index(graph, edge_index_type, dask_client):
             G[et][0] = dask_cudf.from_cudf(cudf.Series(G[et][0]), npartitions=1)
             G[et][1] = dask_cudf.from_cudf(cudf.Series(G[et][1]), npartitions=1)
 
-    cugraph_store = CuGraphStore(F, G, N, multi_gpu=True)
+    cugraph_store = CuGraphStore(F, G, N, multi_gpu=True, order="CSC")
 
     for pyg_can_edge_type in G:
         src, dst = cugraph_store.get_edge_index(
@@ -386,3 +386,29 @@ def test_mg_frame_handle(graph, dask_client):
     F, G, N = graph
     cugraph_store = CuGraphStore(F, G, N, multi_gpu=True)
     assert isinstance(cugraph_store._EXPERIMENTAL__CuGraphStore__graph._plc_graph, dict)
+
+
+@pytest.mark.skipif(isinstance(torch, MissingModule), reason="torch not available")
+def test_cugraph_loader_large_index(dask_client):
+    large_index = (
+        np.random.randint(0, 1_000_000, (100_000_000,)),
+        np.random.randint(0, 1_000_000, (100_000_000,)),
+    )
+
+    large_features = np.random.randint(0, 50, (1_000_000,))
+    F = cugraph.gnn.FeatureStore(backend="torch")
+    F.add_data(large_features, "N", "f")
+
+    store = CuGraphStore(
+        F,
+        {("N", "e", "N"): large_index},
+        {"N": 1_000_000},
+        multi_gpu=True,
+    )
+
+    graph = store._subgraph()
+    assert isinstance(graph, cugraph.Graph)
+
+    el = graph.view_edge_list().compute()
+    assert (el["src"].values_host - large_index[0]).sum() == 0
+    assert (el["dst"].values_host - large_index[1]).sum() == 0
diff --git a/python/cugraph-pyg/cugraph_pyg/tests/test_cugraph_store.py b/python/cugraph-pyg/cugraph_pyg/tests/test_cugraph_store.py
@@ -113,7 +113,7 @@ def test_get_edge_index(graph, edge_index_type):
             G[et][0] = cudf.Series(G[et][0])
             G[et][1] = cudf.Series(G[et][1])
 
-    cugraph_store = CuGraphStore(F, G, N)
+    cugraph_store = CuGraphStore(F, G, N, order="CSC")
 
     for pyg_can_edge_type in G:
         src, dst = cugraph_store.get_edge_index(