create unique token and delete ddf after compute call

rapidsai · Sep 27, 2023 · 4f588de · 4f588de
1 parent dde98b3
commit 4f588de
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/python/cugraph/cugraph/structure/graph_implementation/simpleDistributedGraph.py b/python/cugraph/cugraph/structure/graph_implementation/simpleDistributedGraph.py
@@ -343,6 +343,7 @@ def __from_edgelist(
         }
         wait(list(self._plc_graph.values()))
         del delayed_tasks_d
+        del ddf
         _client.run(gc.collect)
 
     @property
@@ -1195,5 +1196,7 @@ def _get_column_from_ls_dfs(lst_df, col_name):
     if len_df == 0:
         return lst_df[0][col_name]
     output_col = cudf.concat([df[col_name] for df in lst_df], ignore_index=True)
-    # FIXME: For now, don't delete the copied dataframe to avoid crash
+    for df in lst_df:
+        df.drop(columns=[col_name], inplace=True)
+    gc.collect()
     return output_col