[BUG] Fix Graph Construction From Pandas in cuGraph-PyG

dependencies.yaml

@@ -503,9 +503,9 @@ dependencies:
       - output_types: [conda]
         packages:
           - cugraph==23.12.*
-          - pytorch==2.0
+          - pytorch==2.1
           - pytorch-cuda==11.8
-          - pyg=2.3.1=*torch_2.0.0*cu118*
+          - pyg=2.4.0=*torch_2.1.0*cu118*
 
   depends_on_rmm:
     common:

python/cugraph-pyg/conda/cugraph_pyg_dev_cuda-118.yaml

@@ -13,13 +13,13 @@ dependencies:
 - cugraph==23.12.*
 - pandas
 - pre-commit
-- pyg=2.3.1=*torch_2.0.0*cu118*
+- pyg=2.4.0=*torch_2.1.0*cu118*
 - pylibcugraphops==23.12.*
 - pytest
 - pytest-benchmark
 - pytest-cov
 - pytest-xdist
 - pytorch-cuda==11.8
-- pytorch==2.0
+- pytorch==2.1
 - scipy
 name: cugraph_pyg_dev_cuda-118

python/cugraph-pyg/cugraph_pyg/data/cugraph_store.py

@@ -27,11 +27,12 @@
 import cugraph
 import warnings
 
-from cugraph.utilities.utils import import_optional, MissingModule
+import dask.array as dar
+import dask.dataframe as dd
+import dask.distributed as distributed
+import dask_cudf
 
-dd = import_optional("dask.dataframe")
-distributed = import_optional("dask.distributed")
-dask_cudf = import_optional("dask_cudf")
+from cugraph.utilities.utils import import_optional, MissingModule
 
 torch = import_optional("torch")
 torch_geometric = import_optional("torch_geometric")
@@ -364,6 +365,24 @@ def __infer_offsets(
             }
         )
 
+    def __dask_array_from_numpy(
+        self, array: np.ndarray, client: distributed.client.Client, npartitions: int
+    ):
+        split_array = np.array_split(array, npartitions)
+        shapes = [n.shape for n in split_array]
+        scattered_array = client.scatter(split_array)
+        del split_array
+
+        dask_array = dar.concatenate(
+            [
+                dar.from_delayed(s, shape=shapes[i], dtype=array.dtype)
+                for i, s in enumerate(scattered_array)
+            ]
+        )
+        del scattered_array
+
+        return dask_array
+
     def __construct_graph(
         self,
         edge_info: Dict[Tuple[str, str, str], List[TensorType]],
@@ -461,32 +480,42 @@ def __construct_graph(
             ]
         )
 
-        df = pandas.DataFrame(
-            {
-                "src": pandas.Series(na_dst)
-                if order == "CSC"
-                else pandas.Series(na_src),
-                "dst": pandas.Series(na_src)
-                if order == "CSC"
-                else pandas.Series(na_dst),
-                "etp": pandas.Series(na_etp),
-            }
-        )
-        vertex_dtype = df.src.dtype
+        vertex_dtype = na_src.dtype
 
         if multi_gpu:
-            nworkers = len(distributed.get_client().scheduler_info()["workers"])
-            df = dd.from_pandas(df, npartitions=nworkers if len(df) > 32 else 1)
+            client = distributed.get_client()
+            nworkers = len(client.scheduler_info()["workers"])
+            npartitions = nworkers * 4
+
+            src_dar = self.__dask_array_from_numpy(na_src, client, npartitions)
+            del na_src
+
+            dst_dar = self.__dask_array_from_numpy(na_dst, client, npartitions)
+            del na_dst
+
+            etp_dar = self.__dask_array_from_numpy(na_etp, client, npartitions)
+            del na_etp
+
+            df = dd.from_dask_array(etp_dar, columns=["etp"])
+            df["src"] = dst_dar if order == "CSC" else src_dar
+            df["dst"] = src_dar if order == "CSC" else dst_dar
+
+            del src_dar
+            del dst_dar
+            del etp_dar
+
+            if df.etp.dtype != "int32":
+                raise ValueError("Edge type must be int32!")
 
             # Ensure the dataframe is constructed on each partition
             # instead of adding additional synchronization head from potential
             # host to device copies.
             def get_empty_df():
                 return cudf.DataFrame(
                     {
+                        "etp": cudf.Series([], dtype="int32"),
                         "src": cudf.Series([], dtype=vertex_dtype),
                         "dst": cudf.Series([], dtype=vertex_dtype),
-                        "etp": cudf.Series([], dtype="int32"),
                     }
                 )
 
@@ -497,9 +526,23 @@ def get_empty_df():
                 if len(f) > 0
                 else get_empty_df(),
                 meta=get_empty_df(),
-            ).reset_index(drop=True)
+            ).reset_index(
+                drop=True
+            )  # should be ok for dask
         else:
-            df = cudf.from_pandas(df).reset_index(drop=True)
+            df = pandas.DataFrame(
+                {
+                    "src": pandas.Series(na_dst)
+                    if order == "CSC"
+                    else pandas.Series(na_src),
+                    "dst": pandas.Series(na_src)
+                    if order == "CSC"
+                    else pandas.Series(na_dst),
+                    "etp": pandas.Series(na_etp),
+                }
+            )
+            df = cudf.from_pandas(df)
+            df.reset_index(drop=True, inplace=True)
 
         graph = cugraph.MultiGraph(directed=True)
         if multi_gpu:
@@ -518,6 +561,7 @@ def get_empty_df():
                 edge_type="etp",
             )
 
+        del df
         return graph
 
     @property

python/cugraph-pyg/cugraph_pyg/loader/cugraph_node_loader.py

@@ -15,6 +15,7 @@
 
 import os
 import re
+import warnings
 
 import cupy
 import cudf
@@ -159,23 +160,34 @@ def __init__(
         if batch_size is None or batch_size < 1:
             raise ValueError("Batch size must be >= 1")
 
-        self.__directory = tempfile.TemporaryDirectory(dir=directory)
+        self.__directory = (
+            tempfile.TemporaryDirectory() if directory is None else directory
+        )
 
         if isinstance(num_neighbors, dict):
             raise ValueError("num_neighbors dict is currently unsupported!")
 
-        renumber = (
-            True
-            if (
-                (len(self.__graph_store.node_types) == 1)
-                and (len(self.__graph_store.edge_types) == 1)
+        if "renumber" in kwargs:
+            warnings.warn(
+                "Setting renumbering manually could result in invalid output,"
+                " please ensure you intended to do this."
+            )
+            renumber = kwargs.pop("renumber")
+        else:
+            renumber = (
+                True
+                if (
+                    (len(self.__graph_store.node_types) == 1)
+                    and (len(self.__graph_store.edge_types) == 1)
+                )
+                else False
             )
-            else False
-        )
 
         bulk_sampler = BulkSampler(
             batch_size,
-            self.__directory.name,
+            self.__directory
+            if isinstance(self.__directory, str)
+            else self.__directory.name,
             self.__graph_store._subgraph(edge_types),
             fanout_vals=num_neighbors,
             with_replacement=replace,
@@ -219,7 +231,13 @@ def __init__(
             )
 
         bulk_sampler.flush()
-        self.__input_files = iter(os.listdir(self.__directory.name))
+        self.__input_files = iter(
+            os.listdir(
+                self.__directory
+                if isinstance(self.__directory, str)
+                else self.__directory.name
+            )
+        )
 
     def __next__(self):
         from time import perf_counter
@@ -437,11 +455,10 @@ def __next__(self):
 
         # Account for CSR format in cuGraph vs. CSC format in PyG
         if self.__coo and self.__graph_store.order == "CSC":
-            for node_type in out.edge_index_dict:
-                out[node_type].edge_index[0], out[node_type].edge_index[1] = (
-                    out[node_type].edge_index[1],
-                    out[node_type].edge_index[0],
-                )
+            for edge_type in out.edge_index_dict:
+                src = out[edge_type].edge_index[0]
+                dst = out[edge_type].edge_index[1]
+                out[edge_type].edge_index = torch.stack([dst, src])
 
         out.set_value_dict("num_sampled_nodes", sampler_output.num_sampled_nodes)
         out.set_value_dict("num_sampled_edges", sampler_output.num_sampled_edges)

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")

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

python/cugraph-pyg/cugraph_pyg/tests/test_cugraph_loader.py

@@ -18,6 +18,7 @@
 
 import cudf
 import cupy
+import numpy as np
 
 from cugraph_pyg.loader import CuGraphNeighborLoader
 from cugraph_pyg.loader import BulkSampleLoader
@@ -27,6 +28,8 @@
 from cugraph.gnn import FeatureStore
 from cugraph.utilities.utils import import_optional, MissingModule
 
+from typing import Dict, Tuple
+
 torch = import_optional("torch")
 torch_geometric = import_optional("torch_geometric")
 trim_to_layer = import_optional("torch_geometric.utils.trim_to_layer")
@@ -40,7 +43,9 @@
 
 
 @pytest.mark.skipif(isinstance(torch, MissingModule), reason="torch not available")
-def test_cugraph_loader_basic(karate_gnn):
+def test_cugraph_loader_basic(
+    karate_gnn: Tuple[FeatureStore, Dict[str, np.ndarray], Dict[str, int]]
+):
     F, G, N = karate_gnn
     cugraph_store = CuGraphStore(F, G, N, order="CSR")
     loader = CuGraphNeighborLoader(
@@ -66,7 +71,9 @@ def test_cugraph_loader_basic(karate_gnn):
 
 
 @pytest.mark.skipif(isinstance(torch, MissingModule), reason="torch not available")
-def test_cugraph_loader_hetero(karate_gnn):
+def test_cugraph_loader_hetero(
+    karate_gnn: Tuple[FeatureStore, Dict[str, np.ndarray], Dict[str, int]]
+):
     F, G, N = karate_gnn
     cugraph_store = CuGraphStore(F, G, N, order="CSR")
     loader = CuGraphNeighborLoader(
@@ -342,7 +349,7 @@ def test_cugraph_loader_e2e_coo():
 @pytest.mark.skipif(isinstance(torch, MissingModule), reason="torch not available")
 @pytest.mark.skipif(not HAS_TORCH_SPARSE, reason="torch-sparse not available")
 @pytest.mark.parametrize("framework", ["pyg", "cugraph-ops"])
-def test_cugraph_loader_e2e_csc(framework):
+def test_cugraph_loader_e2e_csc(framework: str):
     m = [2, 9, 99, 82, 9, 3, 18, 1, 12]
     x = torch.randint(3000, (256, 256)).to(torch.float32)
     F = FeatureStore()
@@ -442,3 +449,38 @@ def test_cugraph_loader_e2e_csc(framework):
         x = x.narrow(dim=0, start=0, length=s - num_sampled_nodes[1])
 
         assert list(x.shape) == [1, 1]
+
+
+@pytest.mark.skipif(isinstance(torch, MissingModule), reason="torch not available")
+@pytest.mark.parametrize("directory", ["local", "temp"])
+def test_load_directory(
+    karate_gnn: Tuple[FeatureStore, Dict[str, np.ndarray], Dict[str, int]],
+    directory: str,
+):
+    if directory == "local":
+        local_dir = tempfile.TemporaryDirectory(dir=".")
+
+    cugraph_store = CuGraphStore(*karate_gnn)
+    cugraph_loader = CuGraphNeighborLoader(
+        (cugraph_store, cugraph_store),
+        torch.arange(8, dtype=torch.int64),
+        2,
+        num_neighbors=[8, 4, 2],
+        random_state=62,
+        replace=False,
+        directory=None if directory == "temp" else local_dir.name,
+        batches_per_partition=1,
+    )
+
+    it = iter(cugraph_loader)
+    next_batch = next(it)
+    assert next_batch is not None
+
+    if directory == "local":
+        assert len(os.listdir(local_dir.name)) == 4
+
+    count = 1
+    while next(it, None) is not None:
+        count += 1
+
+    assert count == 4

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(