Ensure correct data type

python/cugraph/cugraph/dask/traversal/bfs.py

@@ -35,15 +35,23 @@ def convert_to_cudf(cp_arrays):
  return df
 
 
-def _call_plc_bfs(sID, mg_graph_x, st_x, depth_limit=None, return_distances=True):
+def _call_plc_bfs(
+ sID,
+ mg_graph_x,
+ st_x,
+ depth_limit=None,
+ direction_optimizing=False,
+ return_distances=True,
+ do_expensive_check=False,
+):
  return pylibcugraph_bfs(
  ResourceHandle(Comms.get_handle(sID).getHandle()),
- mg_graph_x,
- cudf.Series(st_x, dtype="int32"),
- False,
- depth_limit if depth_limit is not None else 0,
- return_distances,
- True,
+ graph=mg_graph_x,
+ sources=st_x,
+ direction_optimizing=direction_optimizing,
+ depth_limit=depth_limit if depth_limit is not None else 0,
+ compute_predecessors=return_distances,
+ do_expensive_check=do_expensive_check,
  )
 
 
@@ -156,6 +164,10 @@ def bfs(input_graph, start, depth_limit=None, return_distances=True, check_start
  start = input_graph.lookup_internal_vertex_id(start, tmp_col_names)
 
  data_start = get_distributed_data(start)
+ do_expensive_check = False
+ # FIXME: Why is 'direction_optimizing' not part of the python cugraph API
+ # and why is it set to 'False' by default
+ direction_optimizing = False
 
  cupy_result = [
  client.submit(
@@ -164,7 +176,9 @@ def bfs(input_graph, start, depth_limit=None, return_distances=True, check_start
  input_graph._plc_graph[w],
  st[0],
  depth_limit,
+ direction_optimizing,
  return_distances,
+ do_expensive_check,
  workers=[w],
  allow_other_workers=False,
  )

python/cugraph/cugraph/structure/number_map.py

@@ -258,15 +258,18 @@ def indirection_map(self, ddf, src_col_names, dst_col_names):
 
  # Set global index
  tmp_ddf = tmp_ddf.assign(idx=1)
- tmp_ddf["global_id"] = tmp_ddf.idx.cumsum() - 1
+ # ensure the original vertex and the 'global_id' columns are
+ # of the same type unless the original vertex type is 'string'
+ tmp_ddf["global_id"] = tmp_ddf.idx.cumsum().astype(self.id_type) - 1
  tmp_ddf = tmp_ddf.drop(columns="idx")
  tmp_ddf = tmp_ddf.persist()
  self.ddf = tmp_ddf
  return tmp_ddf
 
- def __init__(self, id_type=np.int32, renumber_type=None):
+ def __init__(self, renumber_id_type=np.int32, unrenumbered_id_type=np.int32):
  self.implementation = None
- self.id_type = id_type
+ self.renumber_id_type = renumber_id_type
+ self.unrenumbered_id_type = unrenumbered_id_type
  # The default src/dst column names in the resulting renumbered
  # dataframe. These may be updated by the renumbering methods if the
  # input dataframe uses the default names.
@@ -479,13 +482,29 @@ def renumber_and_segment(
  legacy_renum_only=False,
  ):
  renumbered = True
+
+ # For columns with mismatch dtypes, set the renumbered
+ # id_type to either 'int32' or 'int64'
+ if df.dtypes.nunique() > 1:
+ # can't determine the edgelist input type
+ unrenumbered_id_type = None
+ else:
+ unrenumbered_id_type = df.dtypes[0]
+
+ if np.int64 in list(df.dtypes):
+ renumber_id_type = np.int64
+ else:
+ # renumber the edgelist to 'int32'
+ renumber_id_type = np.int32
+
  # FIXME: Drop the renumber_type 'experimental' once all the
  # algos follow the C/Pylibcugraph path
 
  # The renumber_type 'legacy' runs both the python and the
  # C++ renumbering.
  if isinstance(src_col_names, list):
  renumber_type = "legacy"
+
  elif not (
  df[src_col_names].dtype == np.int32 or df[src_col_names].dtype == np.int64
  ):
@@ -500,7 +519,7 @@ def renumber_and_segment(
  renumber_type = "skip_renumbering"
  renumbered = False
 
- renumber_map = NumberMap()
+ renumber_map = NumberMap(renumber_id_type, unrenumbered_id_type)
  if not isinstance(src_col_names, list):
  src_col_names = [src_col_names]
  dst_col_names = [dst_col_names]
@@ -512,11 +531,19 @@ def renumber_and_segment(
 
  if isinstance(df, cudf.DataFrame):
  renumber_map.implementation = NumberMap.SingleGPU(
- df, src_col_names, dst_col_names, renumber_map.id_type, store_transposed
+ df,
+ src_col_names,
+ dst_col_names,
+ renumber_map.renumber_id_type,
+ store_transposed,
  )
  elif isinstance(df, dask_cudf.DataFrame):
  renumber_map.implementation = NumberMap.MultiGPU(
- df, src_col_names, dst_col_names, renumber_map.id_type, store_transposed
+ df,
+ src_col_names,
+ dst_col_names,
+ renumber_map.renumber_id_type,
+ store_transposed,
  )
  else:
  raise TypeError("df must be cudf.DataFrame or dask_cudf.DataFrame")

python/cugraph/cugraph/tests/mg/test_mg_renumber.py

@@ -308,7 +308,7 @@ def test_pagerank_string_vertex_ids(dask_client):
  G.from_cudf_edgelist(df, source="src", destination="dst")
 
  sg_results = cugraph.pagerank(G)
- sg_results = sg_results.sort_values("pagerank").reset_index(drop=True)
+ sg_results = sg_results.sort_values("vertex").reset_index(drop=True)
 
  # MG
  ddf = dask_cudf.from_cudf(df, npartitions=2)
@@ -318,7 +318,28 @@ def test_pagerank_string_vertex_ids(dask_client):
  mg_results = dcg.pagerank(G_dask)
  # Organize results for easy comparison, this does not change the values. MG
  # Pagerank defaults to float64, so convert to float32 when comparing to SG
- mg_results = mg_results.compute().sort_values("pagerank").reset_index(drop=True)
+ mg_results = mg_results.compute().sort_values("vertex").reset_index(drop=True)
  mg_results["pagerank"] = mg_results["pagerank"].astype("float32")
 
  assert_frame_equal(sg_results, mg_results)
+
+
+@pytest.mark.parametrize("dtype", ["int32", "int64"])
+def test_mg_renumber_multi_column(dtype, dask_client):
+ df = cudf.DataFrame(
+ {"src_a": [i for i in range(0, 10)], "dst_a": [i for i in range(10, 20)]}
+ ).astype(dtype)
+
+ df["src_b"] = df["src_a"] + 10
+ df["dst_b"] = df["dst_a"] + 20
+ src_col = ["src_a", "src_b"]
+ dst_col = ["dst_a", "dst_b"]
+
+ ddf = dask_cudf.from_cudf(df, npartitions=2)
+ edgelist_type = list(ddf.dtypes)
+ G = cugraph.Graph()
+ G.from_dask_cudf_edgelist(ddf, source=src_col, destination=dst_col)
+ renumbered_ddf = G.edgelist.edgelist_df
+ renumbered_edgelist_type = list(renumbered_ddf.dtypes)
+
+ assert set(renumbered_edgelist_type).issubset(set(edgelist_type))