def _convert_df_to_output_type(df, input_type, return_predecessors):
"""
Given a cudf.DataFrame df, convert it to a new type appropriate for the
graph algos in this module, based on input_type.
return_predecessors is only used for return values from cupy/scipy input
types.
"""
if input_type in [Graph, DiGraph, MultiGraph, MultiDiGraph]:
return df
elif is_nx_graph_type(input_type):
return df.to_pandas()
elif is_matrix_type(input_type):
# A CuPy/SciPy input means the return value will be a 2-tuple of:
# distance: cupy.ndarray
# predecessor: cupy.ndarray
sorted_df = df.sort_values("vertex")
if return_predecessors:
if is_cp_matrix_type(input_type):
return (cp.fromDlpack(sorted_df["distance"].to_dlpack()),
cp.fromDlpack(sorted_df["predecessor"].to_dlpack()))
else:
return (sorted_df["distance"].to_numpy(),
sorted_df["predecessor"].to_numpy())
else:
if is_cp_matrix_type(input_type):
return cp.fromDlpack(sorted_df["distance"].to_dlpack())
else:
return sorted_df["distance"].to_numpy()
else:
raise TypeError(f"input type {input_type} is not a supported type.")
def _convert_df_to_output_type(df, input_type):
"""
Given a cudf.DataFrame df, convert it to a new type appropriate for the
graph algos in this module, based on input_type.
"""
if input_type in [Graph, DiGraph]:
return df
elif (nx is not None) and (input_type in [nx.Graph, nx.DiGraph]):
return df.to_pandas()
elif is_matrix_type(input_type):
# A CuPy/SciPy input means the return value will be a 2-tuple of:
# distance: cupy.ndarray
# predecessor: cupy.ndarray
sorted_df = df.sort_values("vertex")
if is_cp_matrix_type(input_type):
distances = cp.fromDlpack(sorted_df["distance"].to_dlpack())
preds = cp.fromDlpack(sorted_df["predecessor"].to_dlpack())
if "sp_counter" in df.columns:
return (distances, preds,
cp.fromDlpack(sorted_df["sp_counter"].to_dlpack()))
else:
return (distances, preds)
else:
distances = sorted_df["distance"].to_array()
preds = sorted_df["predecessor"].to_array()
if "sp_counter" in df.columns:
return (distances, preds, sorted_df["sp_counter"].to_array())
else:
return (distances, preds)
else:
raise TypeError(f"input type {input_type} is not a supported type.")
def _convert_df_to_output_type(df, input_type, return_labels):
"""
Given a cudf.DataFrame df, convert it to a new type appropriate for the
graph algos in this module, based on input_type.
return_labels is only used for return values from cupy/scipy input types.
"""
if input_type in [Graph, DiGraph]:
return df
elif is_nx_graph_type(input_type):
return df_score_to_dictionary(df, "labels", "vertex")
elif is_matrix_type(input_type):
# Convert DF of 2 columns (labels, vertices) to the SciPy-style return
# value:
# n_components: int
# The number of connected components (number of unique labels).
# labels: ndarray
# The length-N array of labels of the connected components.
n_components = len(df["labels"].unique())
sorted_df = df.sort_values("vertex")
if return_labels:
if is_cp_matrix_type(input_type):
labels = cp.fromDlpack(sorted_df["labels"].to_dlpack())
else:
labels = sorted_df["labels"].to_numpy()
return (n_components, labels)
else:
return n_components
else:
raise TypeError(f"input type {input_type} is not a supported type.")
