diff --git a/graphistry/plugins/cugraph.py b/graphistry/plugins/cugraph.py index 20409bc831..993464454e 100644 --- a/graphistry/plugins/cugraph.py +++ b/graphistry/plugins/cugraph.py @@ -1,15 +1,19 @@ -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Optional, Set, Union + +import numpy as np import pandas as pd import warnings from graphistry.constants import NODE +from graphistry.compute.typing import DataFrameT, SeriesT +from graphistry.compute.typing import NodeId from graphistry.Engine import EngineAbstract +from graphistry.Engine import is_cudf_series from graphistry.Plottable import Plottable from graphistry.plugins_types.cugraph_types import CuGraphKind from graphistry.util import setup_logger logger = setup_logger(__name__) - #import logging #logger.setLevel(logging.DEBUG) @@ -30,6 +34,142 @@ def df_to_gdf(df: Any): raise ValueError('Unsupported data type, expected pd.DataFrame or cudf.DataFrame, got: %s', type(df)) +def _unique_ids_to_set(values: Union[SeriesT, np.ndarray]) -> Set[NodeId]: + if is_cudf_series(values): + return set(values.to_arrow().to_pylist()) + if isinstance(values, (pd.Series, np.ndarray)): + return set(values.tolist()) + raise TypeError(f'Unsupported unique-ID type: {type(values)}') + + +def _restore_graph_edge_attrs( + returned_edges: DataFrameT, + original_edges: DataFrameT, + source: str, + destination: str, +) -> DataFrameT: + """Attach original edge attributes to graph-algorithm output without join fanout. + + cuGraph returns one row per endpoint pair. A duplicate original pair cannot be + assigned deterministically, so reject it instead of silently multiplying rows. + """ + try: + import cudf + except ImportError: + cudf = None + original_columns = list(original_edges.columns) + attr_columns = [column for column in original_columns if column not in (source, destination)] + if not attr_columns: + return returned_edges + + returned_columns = list(returned_edges.columns) + overlapping = [column for column in attr_columns if column in returned_columns] + if overlapping: + raise ValueError(f'Returned edge columns overlap original attributes: {overlapping}') + + existing_columns = set(returned_columns) | set(original_columns) + if original_edges.duplicated(subset=[source, destination]).any(): + raise ValueError( + 'cuGraph graph algorithms require unique (source, destination) edge pairs; ' + 'drop or aggregate parallel edges before calling this method' + ) + + returned_id = '__cugraph_returned_edge_id__' + original_id = '__cugraph_original_edge_id__' + forward_match = '__cugraph_forward_match__' + while ( + returned_id in existing_columns + or original_id in existing_columns + or returned_id == original_id + ): + returned_id = f'_{returned_id}' + original_id = f'_{original_id}' + while forward_match in existing_columns or forward_match == returned_id: + forward_match = f'_{forward_match}' + + returned = returned_edges.copy() + original = ( + df_to_gdf(original_edges).copy() + if cudf is not None and isinstance(returned_edges, cudf.DataFrame) + else original_edges.copy() + ) + returned[returned_id] = range(len(returned)) + original[original_id] = range(len(original)) + + forward = returned[[returned_id, source, destination]].merge( + original[[original_id, source, destination, *attr_columns]], + how='inner', + on=[source, destination], + ) + forward[forward_match] = True + reversed_original = original[[original_id, source, destination, *attr_columns]].rename( + columns={source: destination, destination: source} + ) + reversed_matches = returned[[returned_id, source, destination]].merge( + reversed_original, + how='inner', + on=[source, destination], + ) + reversed_matches[forward_match] = False + concat = cudf.concat if cudf is not None and isinstance(returned_edges, cudf.DataFrame) else pd.concat + selected = concat([forward, reversed_matches]).drop_duplicates( + subset=[returned_id, *attr_columns], keep='first' + ) + conflicted_ids = selected.loc[ + selected.duplicated(subset=[returned_id], keep=False), returned_id + ].unique() + forward_rows = returned[[returned_id, source, destination]] + reverse_rows = forward_rows.rename(columns={source: destination, destination: source}) + counterparts = forward_rows.merge( + reverse_rows, on=[source, destination], suffixes=('_left', '_right') + ) + counterpart_ids = counterparts.loc[ + counterparts[f'{returned_id}_left'] != counterparts[f'{returned_id}_right'], + f'{returned_id}_left', + ].unique() + conflicted_id_set = _unique_ids_to_set(conflicted_ids) + counterpart_id_set = _unique_ids_to_set(counterpart_ids) + if conflicted_id_set - counterpart_id_set: + raise ValueError( + 'cuGraph undirected graph algorithms return one orientation per endpoint pair; ' + 'antiparallel original edges with conflicting attributes are ambiguous; ' + 'normalize attributes or use a directed graph' + ) + selected = selected.loc[ + ~selected[returned_id].isin(conflicted_id_set) + | selected[forward_match] + ] + selected = selected.sort_values(returned_id) + + restored = returned.merge( + selected[[returned_id, *attr_columns]], + how='left', + on=[returned_id], + ) + return restored.drop(columns=[returned_id]) + + +def _restore_graph_return_edge_attrs( + graph: Plottable, + original_edges: Optional[DataFrameT], + source: Optional[str], + destination: Optional[str], +) -> Plottable: + """Restore caller edge bindings and attributes on a returned cuGraph graph.""" + assert source is not None and destination is not None + if original_edges is not None: + attr_columns = [ + column for column in original_edges.columns + if column not in (source, destination) + ] + if attr_columns: + restored = _restore_graph_edge_attrs( + graph._edges, original_edges, source, destination + ) + return graph.edges(restored, source, destination) + return graph + + def from_cugraph(self, G, node_attributes: Optional[List[str]] = None, @@ -357,15 +497,11 @@ def compute_cugraph_core( out = out[0] if out_col is not None: raise ValueError('Graph returned, but out_col was specified') - self2 = self - if self._source != out.source_columns: - logger.debug('Switching g source column name to G source column name to work around cugraph inconsistency') - if out.source_columns == 'src': - self2 = self.edges(self._edges.rename(columns={self._source: 'src', self._destination: 'dst'}), 'src', 'dst') - res = from_cugraph(self2, out, load_nodes=False) - if not (self2 is self): - res = res.edges(self._edges, self._source, self._destination) - return res + if out.source_columns == 'src' and out.destination_columns == 'dst': + out.source_columns = self._source + out.destination_columns = self._destination + res = from_cugraph(self, out, load_nodes=False, merge_if_existing=False) + return _restore_graph_return_edge_attrs(res, self._edges, self._source, self._destination) raise ValueError('Unsupported algorithm: %s', alg) diff --git a/graphistry/tests/plugins/test_cugraph.py b/graphistry/tests/plugins/test_cugraph.py index 883ac452e0..3793e7a3cf 100644 --- a/graphistry/tests/plugins/test_cugraph.py +++ b/graphistry/tests/plugins/test_cugraph.py @@ -1,9 +1,11 @@ from functools import lru_cache +from collections import Counter import graphistry, logging, os, pandas as pd, pytest from graphistry.tests.common import NoAuthTestCase from graphistry.constants import SRC, DST, NODE from graphistry.plugins.cugraph import ( SRC_CUGRAPH, DST_CUGRAPH, NODE_CUGRAPH, + _restore_graph_edge_attrs, _restore_graph_return_edge_attrs, from_cugraph, to_cugraph, compute_cugraph, layout_cugraph, compute_algs, @@ -470,6 +472,141 @@ def test_enrich_with_stat_direct(self): @pytest.mark.skipif(not test_cugraph, reason="Requires TEST_CUGRAPH=1") class Test_cugraph_compute(NoAuthTestCase): + def test_graph_return_endpoint_column_collision(self): + import cudf + + for collision_columns in (("src",), ("dst",), ("src", "dst")): + with self.subTest(collision_columns=collision_columns): + data = {"s": ["a", "b", "c"], "d": ["b", "c", "a"]} + for column in collision_columns: + data[column] = [f"{column}_0", f"{column}_1", f"{column}_2"] + g = graphistry.edges(cudf.DataFrame(data), "s", "d") + + out1 = g.compute_cugraph("k_core", directed=False, params={"k": 2}) + out2 = out1.compute_cugraph("k_core", directed=False, params={"k": 2}) + + assert out2._source == "s" and out2._destination == "d" + result_edges = Counter( + frozenset(edge) for edge in zip(out2._edges["s"].to_pandas(), out2._edges["d"].to_pandas()) + ) + expected_edges = Counter(frozenset(edge) for edge in zip(data["s"], data["d"])) + assert result_edges == expected_edges + for column in collision_columns: + assert sorted(out2._edges[column].to_pandas()) == sorted(data[column]) + + def test_graph_return_parallel_edges_declined(self): + import cudf + + data = { + "s": ["a", "a", "b"], + "d": ["b", "b", "c"], + "label": ["parallel_0", "parallel_1", "bc"], + } + g = graphistry.edges(cudf.DataFrame(data), "s", "d") + + with pytest.raises(ValueError, match="unique \\(source, destination\\) edge pairs"): + g.compute_cugraph("k_core", directed=False, params={"k": 2}) + + def test_graph_return_clean_control(self): + import cudf + + data = { + "s": ["a", "b", "c"], + "d": ["b", "c", "a"], + "source": ["source_0", "source_1", "source_2"], + "target": ["target_0", "target_1", "target_2"], + "__attr_source_2__": ["generated_0", "generated_1", "generated_2"], + } + g = graphistry.edges(cudf.DataFrame(data), "s", "d") + + out = g.compute_cugraph("k_core", directed=False, params={"k": 2}) + + assert out._source == "s" and out._destination == "d" + expected_attrs = { + frozenset((data["s"][i], data["d"][i])): ( + data["source"][i], data["target"][i], data["__attr_source_2__"][i] + ) + for i in range(len(data["s"])) + } + result_edges = { + frozenset(edge) for edge in zip(out._edges["s"].to_pandas(), out._edges["d"].to_pandas()) + } + result_rows = list(zip( + out._edges["s"].to_pandas(), + out._edges["d"].to_pandas(), + out._edges["source"].to_pandas(), + out._edges["target"].to_pandas(), + out._edges["__attr_source_2__"].to_pandas(), + )) + result_attrs = { + frozenset((source, destination)): (source_value, target_value, generated_value) + for source, destination, source_value, target_value, generated_value in result_rows + } + assert result_attrs == { + edge: expected_attrs[edge] for edge in expected_attrs + if edge in result_edges + } + + def test_graph_return_directed_ego_preserves_antiparallel_attrs(self): + import cudf + + data = { + "s": ["a", "b", "b", "c"], + "d": ["b", "a", "c", "b"], + "label": ["ab", "ba", "bc", "cb"], + } + g = graphistry.edges(cudf.DataFrame(data), "s", "d") + + out = g.compute_cugraph( + "ego_graph", directed=True, params={"n": cudf.Series(["a"]), "radius": 1} + ) + + assert out._source == "s" and out._destination == "d" + result_rows = list(zip( + out._edges["s"].to_pandas(), + out._edges["d"].to_pandas(), + out._edges["label"].to_pandas(), + )) + assert len(result_rows) == 2 + assert set(result_rows) == {("a", "b", "ab"), ("b", "a", "ba")} + + def test_graph_return_filters_topology_and_preserves_retained_attrs(self): + import cudf + + data = { + "s": ["a", "b", "c", "a"], + "d": ["b", "c", "a", "tail"], + "label": ["triangle_ab", "triangle_bc", "triangle_ca", "tail_at"], + } + g = graphistry.edges(cudf.DataFrame(data), "s", "d") + + out1 = g.compute_cugraph("k_core", directed=False, params={"k": 2}) + out2 = out1.compute_cugraph("k_core", directed=False, params={"k": 2}) + + assert out2._source == "s" and out2._destination == "d" + assert len(out1._edges) == 3 + assert len(out2._edges) == 3 + expected_labels = { + frozenset((data["s"][i], data["d"][i])): data["label"][i] + for i in range(len(data["s"])) + } + result_edges = { + frozenset(edge) for edge in zip(out2._edges["s"].to_pandas(), out2._edges["d"].to_pandas()) + } + result_rows = list(zip( + out2._edges["s"].to_pandas(), + out2._edges["d"].to_pandas(), + out2._edges["label"].to_pandas(), + )) + assert { + frozenset((source, destination)): label + for source, destination, label in result_rows + } == { + edge: expected_labels[edge] for edge in expected_labels + if edge in result_edges + } + + def test_all_calls(self): import cudf, cugraph @@ -624,6 +761,7 @@ def test_all_calls(self): assert len(out._nodes.columns) == len(g._nodes.columns) + len(out_cols) assert out._edges.shape == g._edges.shape + def test_chain(self): import cudf, cugraph @@ -655,6 +793,158 @@ def test_chain(self): assert out._edges.shape == g._edges.shape +class Test_restore_graph_edge_attrs: + def test_preserves_rows_and_attrs_across_directions(self): + returned = pd.DataFrame({"s": ["a", "b"], "d": ["b", "a"]}) + original = pd.DataFrame({ + "s": ["a"], "d": ["b"], "label": ["ab"], "weight": [2.0] + }) + + restored = _restore_graph_edge_attrs(returned, original, "s", "d") + + assert restored.to_dict("records") == [ + {"s": "a", "d": "b", "label": "ab", "weight": 2.0}, + {"s": "b", "d": "a", "label": "ab", "weight": 2.0}, + ] + + def test_selects_orientation_for_opposite_edges(self): + returned = pd.DataFrame({"s": ["a", "b"], "d": ["b", "a"]}) + original = pd.DataFrame({ + "s": ["a", "b"], "d": ["b", "a"], "label": ["ab", "ba"] + }) + + restored = _restore_graph_edge_attrs(returned, original, "s", "d") + + assert list(zip(restored["s"], restored["d"], restored["label"])) == [ + ("a", "b", "ab"), ("b", "a", "ba") + ] + + def test_preserves_null_attrs(self): + returned = pd.DataFrame({"s": ["a"], "d": ["b"]}) + original = pd.DataFrame({"s": ["a"], "d": ["b"], "label": [None]}) + + restored = _restore_graph_edge_attrs(returned, original, "s", "d") + + assert restored["label"].tolist() == [None] + + def test_declines_parallel_edges(self): + returned = pd.DataFrame({"s": ["a"], "d": ["b"]}) + original = pd.DataFrame({ + "s": ["a", "a"], "d": ["b", "b"], "label": ["one", "two"] + }) + + with pytest.raises(ValueError, match="unique \\(source, destination\\) edge pairs"): + _restore_graph_edge_attrs(returned, original, "s", "d") + + def test_declines_antiparallel_attribute_conflicts(self): + returned = pd.DataFrame({"s": ["a"], "d": ["b"]}) + original = pd.DataFrame({ + "s": ["a", "b"], "d": ["b", "a"], "label": ["ab", "ba"], + }) + + with pytest.raises(ValueError, match="antiparallel original edges"): + _restore_graph_edge_attrs(returned, original, "s", "d") + + def test_allows_antiparallel_identical_attributes(self): + returned = pd.DataFrame({"s": ["a"], "d": ["b"]}) + original = pd.DataFrame({ + "s": ["a", "b"], "d": ["b", "a"], "label": ["same", "same"], + }) + + restored = _restore_graph_edge_attrs(returned, original, "s", "d") + + assert restored.shape == (1, 3) + assert restored["label"].tolist() == ["same"] + + def test_returns_edges_without_attributes(self): + returned = pd.DataFrame({"s": ["a"], "d": ["b"]}) + original = pd.DataFrame({"s": ["a"], "d": ["b"]}) + + restored = _restore_graph_edge_attrs(returned, original, "s", "d") + + assert restored.to_dict("records") == [{"s": "a", "d": "b"}] + + def test_declines_returned_attribute_collision(self): + returned = pd.DataFrame({"s": ["a"], "d": ["b"], "label": ["new"]}) + original = pd.DataFrame({"s": ["a"], "d": ["b"], "label": ["old"]}) + + with pytest.raises(ValueError, match="overlap original attributes"): + _restore_graph_edge_attrs(returned, original, "s", "d") + + def test_avoids_reserved_row_identifiers(self): + returned = pd.DataFrame({ + "__cugraph_returned_edge_id__": ["x"], "s": ["a"], "d": ["b"], + }) + original = pd.DataFrame({ + "__cugraph_original_edge_id__": ["y"], "s": ["a"], "d": ["b"], "label": ["ab"], + }) + + restored = _restore_graph_edge_attrs(returned, original, "s", "d") + + assert restored["label"].tolist() == ["ab"] + + def test_avoids_cross_frame_returned_identifier(self): + returned = pd.DataFrame({"s": ["a"], "d": ["b"]}) + original = pd.DataFrame({ + "s": ["a"], "d": ["b"], + "__cugraph_returned_edge_id__": ["user"], "label": ["ab"], + }) + + restored = _restore_graph_edge_attrs(returned, original, "s", "d") + + assert restored.shape == (1, 4) + assert restored[["s", "d"]].to_dict("records") == [{"s": "a", "d": "b"}] + assert restored["__cugraph_returned_edge_id__"].tolist() == ["user"] + assert restored["label"].tolist() == ["ab"] + + def test_avoids_cross_frame_original_identifier(self): + returned = pd.DataFrame({ + "s": ["a"], "d": ["b"], "__cugraph_original_edge_id__": ["returned"], + }) + original = pd.DataFrame({ + "s": ["a"], "d": ["b"], + "__cugraph_returned_edge_id__": ["original"], "label": ["ab"], + }) + + restored = _restore_graph_edge_attrs(returned, original, "s", "d") + + assert restored.shape == (1, 5) + assert restored[["s", "d"]].to_dict("records") == [{"s": "a", "d": "b"}] + assert restored["__cugraph_original_edge_id__"].tolist() == ["returned"] + assert restored["__cugraph_returned_edge_id__"].tolist() == ["original"] + assert restored["label"].tolist() == ["ab"] + + def test_restores_graph_return_attrs(self): + graph = graphistry.edges(pd.DataFrame({"s": ["a", "b"], "d": ["b", "a"]}), "s", "d") + original = pd.DataFrame({"s": ["a"], "d": ["b"], "label": ["ab"]}) + + restored = _restore_graph_return_edge_attrs(graph, original, "s", "d") + + assert restored._source == "s" and restored._destination == "d" + assert restored._edges.sort_values("s")["label"].tolist() == ["ab", "ab"] + + def test_keeps_graph_return_without_attrs(self): + graph = graphistry.edges(pd.DataFrame({"s": ["a"], "d": ["b"]}), "s", "d") + + restored = _restore_graph_return_edge_attrs(graph, None, "s", "d") + + assert restored is graph + + def test_restores_cross_frame_reserved_return_attrs(self): + graph = graphistry.edges(pd.DataFrame({"s": ["a"], "d": ["b"]}), "s", "d") + original = pd.DataFrame({ + "s": ["a"], "d": ["b"], + "__cugraph_returned_edge_id__": ["user"], "label": ["ab"], + }) + + restored = _restore_graph_return_edge_attrs(graph, original, "s", "d") + + assert restored._source == "s" and restored._destination == "d" + assert restored._edges.shape == (1, 4) + assert restored._edges["__cugraph_returned_edge_id__"].tolist() == ["user"] + assert restored._edges["label"].tolist() == ["ab"] + + @pytest.mark.skipif(not test_cugraph, reason="Requires TEST_CUGRAPH=1") class Test_cugraph_layouts(NoAuthTestCase):