From 6c333d3fe50b842f39f92f096212db3b8c74ca8b Mon Sep 17 00:00:00 2001
From: Ralph Liu <137829296+nv-rliu@users.noreply.github.com>
Date: Fri, 24 May 2024 05:13:15 -0400
Subject: [PATCH] Call New `replicate_edgelist` Function (#4441)

Closes #4440

This PR updates `enable_batch` to use the updated implementation for `replicate_edgelist`.

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

Approvers:
  - Rick Ratzel (https://github.com/rlratzel)
  - Joseph Nke (https://github.com/jnke2016)

URL: https://github.com/rapidsai/cugraph/pull/4441
---
 .../cugraph/structure/graph_implementation/simpleGraph.py | 7 ++-----
 python/cugraph/cugraph/structure/replicate_edgelist.py    | 8 ++------
 2 files changed, 4 insertions(+), 11 deletions(-)

diff --git a/python/cugraph/cugraph/structure/graph_implementation/simpleGraph.py b/python/cugraph/cugraph/structure/graph_implementation/simpleGraph.py
index c90607f9bf6..2b974ee3ebf 100644
--- a/python/cugraph/cugraph/structure/graph_implementation/simpleGraph.py
+++ b/python/cugraph/cugraph/structure/graph_implementation/simpleGraph.py
@@ -12,6 +12,7 @@
 # limitations under the License.
 
 from cugraph.structure import graph_primtypes_wrapper
+from cugraph.structure.replicate_edgelist import replicate_cudf_dataframe
 from cugraph.structure.symmetrize import symmetrize
 from cugraph.structure.number_map import NumberMap
 import cugraph.dask.common.mg_utils as mg_utils
@@ -680,16 +681,12 @@ def enable_batch(self):
 
     def _replicate_edgelist(self):
         client = mg_utils.get_client()
-        comms = Comms.get_comms()
 
         # FIXME: There  might be a better way to control it
         if client is None:
             return
-        work_futures = replication.replicate_cudf_dataframe(
-            self.edgelist.edgelist_df, client=client, comms=comms
-        )
 
-        self.batch_edgelists = work_futures
+        self.batch_edgelists = replicate_cudf_dataframe(self.edgelist.edgelist_df)
 
     def _replicate_adjlist(self):
         client = mg_utils.get_client()
diff --git a/python/cugraph/cugraph/structure/replicate_edgelist.py b/python/cugraph/cugraph/structure/replicate_edgelist.py
index d413e50e485..728b247c139 100644
--- a/python/cugraph/cugraph/structure/replicate_edgelist.py
+++ b/python/cugraph/cugraph/structure/replicate_edgelist.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, 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
@@ -269,7 +269,6 @@ def replicate_cudf_dataframe(cudf_dataframe):
         )
 
     _client = default_client()
-
     if not isinstance(cudf_dataframe, dask_cudf.DataFrame):
         if isinstance(cudf_dataframe, cudf.DataFrame):
             df = dask_cudf.from_cudf(
@@ -287,10 +286,7 @@ def replicate_cudf_dataframe(cudf_dataframe):
     df = get_persisted_df_worker_map(df, _client)
 
     ddf = _mg_call_plc_replicate(
-        _client,
-        Comms.get_session_id(),
-        df,
-        "dataframe",
+        _client, Comms.get_session_id(), df, "dataframe", cudf_dataframe.columns
     )
 
     return ddf