[IMP] Sample with Offsets in the Bulk Sampler

cpp/src/c_api/graph_mg.cpp

@@ -338,19 +338,19 @@ extern "C" cugraph_error_code_t cugraph_mg_graph_create(
     weight_type = cugraph_data_type_id_t::FLOAT32;
   }
 
-  CAPI_EXPECTS((edge_type_ids == nullptr) || (p_edge_ids->type_ == edge_type),
+  CAPI_EXPECTS((edge_ids == nullptr) || (p_edge_ids->type_ == edge_type),
                CUGRAPH_INVALID_INPUT,
                "Invalid input arguments: Edge id type must match edge (src/dst) type",
                *error);
 
-  CAPI_EXPECTS((edge_type_ids == nullptr) || (p_edge_type_ids->size_ == p_src->size_),
+  CAPI_EXPECTS((edge_ids == nullptr) || (p_edge_ids->size_ == p_src->size_),
                CUGRAPH_INVALID_INPUT,
-               "Invalid input arguments: src size != edge prop size",
+               "Invalid input arguments: src size != edge id prop size",
                *error);
 
-  CAPI_EXPECTS((edge_ids == nullptr) || (p_edge_ids->size_ == p_src->size_),
+  CAPI_EXPECTS((edge_type_ids == nullptr) || (p_edge_type_ids->size_ == p_src->size_),
                CUGRAPH_INVALID_INPUT,
-               "Invalid input arguments: src size != edge prop size",
+               "Invalid input arguments: src size != edge type prop size",
                *error);
 
   cugraph_data_type_id_t edge_type_id_type;

cpp/src/c_api/graph_sg.cpp

@@ -513,26 +513,19 @@ extern "C" cugraph_error_code_t cugraph_sg_graph_create(
     weight_type = cugraph_data_type_id_t::FLOAT32;
   }
 
-  // FIXME:  The combination of edge_ids != nullptr, edge_type_ids == nullptr
-  //         logically should be valid, but the code will currently break if
-  //         that is that is specified
-  CAPI_EXPECTS(
-    (edge_type_ids == nullptr && edge_ids == nullptr) ||
-      (edge_type_ids != nullptr && edge_ids != nullptr),
-    CUGRAPH_INVALID_INPUT,
-    "Invalid input arguments: either none or both of edge ids and edge types must be provided.",
-    *error);
+  CAPI_EXPECTS((edge_ids == nullptr) || (p_edge_ids->type_ == edge_type),
+               CUGRAPH_INVALID_INPUT,
+               "Invalid input arguments: Edge id type must match edge (src/dst) type",
+               *error);
 
-  CAPI_EXPECTS(
-    (edge_type_ids == nullptr && edge_ids == nullptr) || (p_edge_ids->type_ == edge_type),
-    CUGRAPH_INVALID_INPUT,
-    "Invalid input arguments: Edge id type must match edge (src/dst) type",
-    *error);
+  CAPI_EXPECTS((edge_ids == nullptr) || (p_edge_ids->size_ == p_src->size_),
+               CUGRAPH_INVALID_INPUT,
+               "Invalid input arguments: src size != edge id prop size",
+               *error);
 
-  CAPI_EXPECTS((edge_type_ids == nullptr && edge_ids == nullptr) ||
-                 (p_edge_ids->size_ == p_src->size_ && p_edge_type_ids->size_ == p_dst->size_),
+  CAPI_EXPECTS((edge_type_ids == nullptr) || (p_edge_type_ids->size_ == p_src->size_),
                CUGRAPH_INVALID_INPUT,
-               "Invalid input arguments: src size != edge prop size",
+               "Invalid input arguments: src size != edge type prop size",
                *error);
 
   cugraph_data_type_id_t edge_type_id_type = cugraph_data_type_id_t::INT32;

python/cugraph-dgl/cugraph_dgl/dataloading/dataloader.py

@@ -209,14 +209,12 @@ def __iter__(self):
         output_dir = os.path.join(
             self._sampling_output_dir, "epoch_" + str(self.epoch_number)
         )
-        rank = self._rank
         bs = BulkSampler(
             output_path=output_dir,
             batch_size=self._batch_size,
             graph=self._cugraph_graph,
             batches_per_partition=self._batches_per_partition,
             seeds_per_call=self._seeds_per_call,
-            rank=rank,
             fanout_vals=self.graph_sampler._reversed_fanout_vals,
             with_replacement=self.graph_sampler.replace,
         )
@@ -226,7 +224,6 @@ def __iter__(self):
         batch_df = create_batch_df(self.tensorized_indices_ds)
         bs.add_batches(batch_df, start_col_name="start", batch_col_name="batch_id")
         bs.flush()
-        output_dir = output_dir + f"/rank={rank}/"
         self.cugraph_dgl_dataset.set_input_files(input_directory=output_dir)
         self.epoch_number = self.epoch_number + 1
         return super().__iter__()

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

@@ -420,8 +420,6 @@ def __construct_graph(
             {
                 "src": pandas.Series(na_src),
                 "dst": pandas.Series(na_dst),
-                "w": pandas.Series(np.zeros(len(na_src))),
-                "eid": pandas.Series(np.arange(len(na_src))),
                 "etp": pandas.Series(na_etp),
             }
         )
@@ -441,15 +439,15 @@ def __construct_graph(
                 df,
                 source="src",
                 destination="dst",
-                edge_attr=["w", "eid", "etp"],
+                edge_type="etp",
             )
             distributed.get_client().publish_dataset(cugraph_graph=graph)
         else:
             graph.from_cudf_edgelist(
                 df,
                 source="src",
                 destination="dst",
-                edge_attr=["w", "eid", "etp"],
+                edge_type="etp",
             )
 
         return graph

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

@@ -14,6 +14,7 @@
 import tempfile
 
 import os
+import re
 
 import cupy
 import cudf
@@ -31,6 +32,9 @@
 
 
 class EXPERIMENTAL__BulkSampleLoader:
+
+    __ex_parquet_file = re.compile(r"batch=([0-9]+)\-([0-9]+)\.parquet")
+
     def __init__(
         self,
         feature_store: CuGraphStore,
@@ -40,7 +44,6 @@ def __init__(
         shuffle=False,
         edge_types: Sequence[Tuple[str]] = None,
         directory=None,
-        rank=0,
         starting_batch_id=0,
         batches_per_partition=100,
         # Sampler args
@@ -84,10 +87,6 @@ def __init__(
             The path of the directory to write samples to.
             Defaults to a new generated temporary directory.
 
-        rank: int (optional, default=0)
-            The rank of the current worker.  Should be provided
-            when there are multiple workers.
-
         starting_batch_id: int (optional, default=0)
             The starting id for each batch.  Defaults to 0.
             Generally used when loading previously-sampled
@@ -102,16 +101,16 @@ def __init__(
 
         self.__feature_store = feature_store
         self.__graph_store = graph_store
-        self.__rank = rank
-        self.__next_batch = starting_batch_id
-        self.__end_exclusive = starting_batch_id
+        self.__next_batch = -1
+        self.__end_exclusive = -1
         self.__batches_per_partition = batches_per_partition
         self.__starting_batch_id = starting_batch_id
 
         if isinstance(all_indices, int):
             # Will be loading from disk
             self.__num_batches = all_indices
             self.__directory = directory
+            iter(os.listdir(self.__directory))
             return
 
         if batch_size is None or batch_size < 1:
@@ -123,7 +122,6 @@ def __init__(
             batch_size,
             self.__directory.name,
             self.__graph_store._subgraph(edge_types),
-            rank=rank,
             fanout_vals=num_neighbors,
             with_replacement=replace,
             batches_per_partition=self.__batches_per_partition,
@@ -161,33 +159,36 @@ def __init__(
             )
 
         bulk_sampler.flush()
+        self.__input_files = iter(os.listdir(self.__directory.name))
 
     def __next__(self):
-        # Quit iterating if there are no batches left
-        if self.__next_batch >= self.__num_batches + self.__starting_batch_id:
-            raise StopIteration
-
         # Load the next set of sampling results if necessary
         if self.__next_batch >= self.__end_exclusive:
+            if self.__directory is None:
+                raise StopIteration
+
             # Read the next parquet file into memory
             dir_path = (
                 self.__directory
                 if isinstance(self.__directory, str)
                 else self.__directory.name
             )
-            rank_path = os.path.join(dir_path, f"rank={self.__rank}")
 
-            file_end_batch_incl = min(
-                self.__end_exclusive + self.__batches_per_partition - 1,
-                self.__starting_batch_id + self.__num_batches - 1,
-            )
+            # Will raise StopIteration if there are no files left
+            fname = next(self.__input_files)
+
+            m = self.__ex_parquet_file.match(fname)
+            if m is None:
+                raise ValueError(f"Invalid parquet filename {fname}")
+
+            self.__next_batch, end_inclusive = [int(g) for g in m.groups()]
+            self.__end_exclusive = end_inclusive + 1
+
             parquet_path = os.path.join(
-                rank_path,
-                f"batch={self.__end_exclusive}" f"-{file_end_batch_incl}.parquet",
+                dir_path,
+                fname,
             )
 
-            self.__end_exclusive += self.__batches_per_partition
-
             columns = {
                 "sources": "int64",
                 "destinations": "int64",
@@ -212,6 +213,7 @@ def __next__(self):
         if self.__next_batch >= self.__num_batches + self.__starting_batch_id:
             # Won't delete a non-temp dir (since it would just be deleting a string)
             del self.__directory
+            self.__directory = None
 
         # Get and return the sampled subgraph
         if isinstance(torch_geometric, MissingModule):

python/cugraph/cugraph/community/egonet.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2021-2022, NVIDIA CORPORATION.
+# Copyright (c) 2021-2023, NVIDIA CORPORATION.
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
@@ -132,13 +132,14 @@ def ego_graph(G, n, radius=1, center=True, undirected=None, distance=None):
     df = cudf.DataFrame()
     df["src"] = source
     df["dst"] = destination
-    df["weight"] = weight
+    if weight is not None:
+        df["weight"] = weight
 
     if G.renumbered:
         df, src_names = G.unrenumber(df, "src", get_column_names=True)
         df, dst_names = G.unrenumber(df, "dst", get_column_names=True)
     else:
-        # FIXME: THe original 'src' and 'dst' are not stored in 'simpleGraph'
+        # FIXME: The original 'src' and 'dst' are not stored in 'simpleGraph'
         src_names = "src"
         dst_names = "dst"
 

python/cugraph/cugraph/dask/community/egonet.py

@@ -66,7 +66,10 @@ def convert_to_cudf(cp_arrays):
     df = cudf.DataFrame()
     df["src"] = cp_src
     df["dst"] = cp_dst
-    df["weight"] = cp_weight
+    if cp_weight is None:
+        df["weight"] = None
+    else:
+        df["weight"] = cp_weight
 
     offsets = cudf.Series(cp_offsets)
 

python/cugraph/cugraph/dask/sampling/uniform_neighbor_sample.py

@@ -119,6 +119,13 @@ def convert_to_cudf(cp_arrays, weight_t, with_edge_properties, return_offsets=Fa
         df[edge_type_n] = edge_types
         df[hop_id_n] = hop_ids
 
+        print(
+            f"sources: {sources}\n"
+            f"destinations: {destinations}\n"
+            f"batch: {batch_ids}\n"
+            f"offset: {offsets}\n"
+        )
+
         if return_offsets:
             offsets_df = cudf.DataFrame(
                 {
@@ -141,10 +148,11 @@ def convert_to_cudf(cp_arrays, weight_t, with_edge_properties, return_offsets=Fa
         df[dst_n] = cupy_destinations
         df[indices_n] = cupy_indices
 
-        if weight_t == "int32":
-            df.indices = df.indices.astype("int32")
-        elif weight_t == "int64":
-            df.indices = df.indices.astype("int64")
+        if cupy_indices is not None:
+            if weight_t == "int32":
+                df.indices = df.indices.astype("int32")
+            elif weight_t == "int64":
+                df.indices = df.indices.astype("int64")
 
         return df
 
@@ -296,6 +304,7 @@ def uniform_neighbor_sample(
         List of output GPUs (by rank) corresponding to batch
         id labels in the label list.  Used to assign each batch
         id to a GPU.
+        Must be in ascending order (i.e. [0, 0, 1, 2]).
 
     random_state: int, optional
         Random seed to use when making sampling calls.

python/cugraph/cugraph/dask/traversal/sssp.py

@@ -20,7 +20,6 @@
 import cudf
 import dask_cudf
 from pylibcugraph import sssp as pylibcugraph_sssp, ResourceHandle
-import warnings
 
 
 def _call_plc_sssp(
@@ -102,12 +101,12 @@ def sssp(input_graph, source, cutoff=None, check_source=True):
 
     # FIXME: Implement a better way to check if the graph is weighted similar
     # to 'simpleGraph'
-    if len(input_graph.edgelist.edgelist_df.columns) != 3:
-        warning_msg = (
-            "'SSSP' requires the input graph to be weighted: Unweighted "
-            "graphs will not be supported in the next release."
+    if not input_graph.weighted:
+        err_msg = (
+            "'SSSP' requires the input graph to be weighted."
+            "'BFS' should be used instead of 'SSSP' for unweighted graphs."
         )
-        warnings.warn(warning_msg, PendingDeprecationWarning)
+        raise ValueError(err_msg)
 
     client = default_client()