Merge pull request #198 from neo4j/fix-gds-sampling-setup

Fix bug in heterogenous gds

Author: FlorentinD

changelog.md

@@ -9,6 +9,8 @@
 
 ## Bug fixes
 
+* Fixed a bug with `from_gds` where graphs with different relationship types would fail if they had different properties.
+
 
 ## Improvements
 

examples/gds-example.ipynb

@@ -7,7 +7,6 @@
 
 import pandas as pd
 from graphdatascience import Graph, GraphDataScience
-from pandas import Series
 
 from .pandas import _from_dfs
 from .visualization_graph import VisualizationGraph
@@ -24,22 +23,25 @@ def _fetch_node_dfs(
     }
 
 
-def _fetch_rel_df(gds: GraphDataScience, G: Graph) -> pd.DataFrame:
-    relationship_properties = G.relationship_properties()
-    assert isinstance(relationship_properties, Series)
+def _fetch_rel_dfs(gds: GraphDataScience, G: Graph) -> list[pd.DataFrame]:
+    rel_types = G.relationship_types()
 
-    relationship_properties_per_type = relationship_properties.tolist()
-    property_set: set[str] = set()
-    for props in relationship_properties_per_type:
-        if props:
-            property_set.update(props)
+    rel_props = {rel_type: G.relationship_properties(rel_type) for rel_type in rel_types}
 
-    if len(property_set) > 0:
-        return gds.graph.relationshipProperties.stream(
-            G, relationship_properties=list(property_set), separate_property_columns=True
-        )
+    rel_dfs: list[pd.DataFrame] = []
+    # Have to call per stream per relationship type as there was a bug in GDS < 2.21
+    for rel_type, props in rel_props.items():
+        assert isinstance(props, list)
+        if len(props) > 0:
+            rel_df = gds.graph.relationshipProperties.stream(
+                G, relationship_types=rel_type, relationship_properties=list(props), separate_property_columns=True
+            )
+        else:
+            rel_df = gds.graph.relationships.stream(G, relationship_types=[rel_type])
+
+        rel_dfs.append(rel_df)
 
-    return gds.graph.relationships.stream(G)
+    return rel_dfs
 
 
 def from_gds(
@@ -131,7 +133,7 @@ def from_gds(
             for df in node_dfs.values():
                 df.drop(columns=[property_name], inplace=True)
 
-        rel_df = _fetch_rel_df(gds, G_fetched)
+        rel_dfs = _fetch_rel_dfs(gds, G_fetched)
     finally:
         if G_fetched.name() != G.name():
             G_fetched.drop()
@@ -161,12 +163,13 @@ def from_gds(
     if "caption" not in all_actual_node_properties:
         node_df["caption"] = node_df["labels"].astype(str)
 
-    if "caption" not in rel_df.columns:
-        rel_df["caption"] = rel_df["relationshipType"]
+    for rel_df in rel_dfs:
+        if "caption" not in rel_df.columns:
+            rel_df["caption"] = rel_df["relationshipType"]
 
     try:
         return _from_dfs(
-            node_df, rel_df, node_radius_min_max=node_radius_min_max, rename_properties={"__size": "size"}, dropna=True
+            node_df, rel_dfs, node_radius_min_max=node_radius_min_max, rename_properties={"__size": "size"}, dropna=True
         )
     except ValueError as e:
         err_msg = str(e)

python-wrapper/src/neo4j_viz/gds.py

@@ -163,13 +163,15 @@ def test_from_gds_mocked(mocker: MockerFixture) -> None:
             }
         ),
     }
-    rels = pd.DataFrame(
-        {
-            "sourceNodeId": [0, 1, 2],
-            "targetNodeId": [1, 2, 0],
-            "relationshipType": ["REL", "REL2", "REL"],
-        }
-    )
+    rels = [
+        pd.DataFrame(
+            {
+                "sourceNodeId": [0, 1, 2],
+                "targetNodeId": [1, 2, 0],
+                "relationshipType": ["REL", "REL2", "REL"],
+            }
+        )
+    ]
 
     mocker.patch(
         "graphdatascience.Graph.__init__",
@@ -188,7 +190,7 @@ def test_from_gds_mocked(mocker: MockerFixture) -> None:
     mocker.patch("graphdatascience.Graph.node_count", lambda x: sum(len(df) for df in nodes.values()))
     mocker.patch("graphdatascience.GraphDataScience.__init__", lambda x: None)
     mocker.patch("neo4j_viz.gds._fetch_node_dfs", return_value=nodes)
-    mocker.patch("neo4j_viz.gds._fetch_rel_df", return_value=rels)
+    mocker.patch("neo4j_viz.gds._fetch_rel_dfs", return_value=rels)
 
     gds = GraphDataScience()  # type: ignore[call-arg]
     G = Graph()  # type: ignore[call-arg]
@@ -303,16 +305,24 @@ def test_from_gds_hetero(gds: Any) -> None:
             # No 'component' property
         }
     )
-    rels = pd.DataFrame(
+    X_rels = pd.DataFrame(
+        {
+            "sourceNodeId": [1],
+            "targetNodeId": [3],
+            "weight": [1.5],
+            "relationshipType": ["X"],
+        }
+    )
+    Y_rels = pd.DataFrame(
         {
-            "sourceNodeId": [0, 1],
-            "targetNodeId": [2, 3],
-            "weight": [0.5, 1.5],
-            "relationshipType": ["REL", "REL2"],
+            "sourceNodeId": [0],
+            "targetNodeId": [2],
+            "score": [1],
+            "relationshipType": ["Y"],
         }
     )
 
-    with gds.graph.construct("flo", [A_nodes, B_nodes], rels) as G:
+    with gds.graph.construct("flo", [A_nodes, B_nodes], [X_rels, Y_rels]) as G:
         VG = from_gds(
             gds,
             G,
@@ -333,14 +343,13 @@ def test_from_gds_hetero(gds: Any) -> None:
                     e.source,
                     e.target,
                     e.caption,
-                    e.properties["relationshipType"],
-                    e.properties["weight"],
+                    e.properties,
                 )
                 for e in VG.relationships
             ],
             key=lambda x: x[0],
         )
         assert vg_rels == [
-            (0, 2, "REL", "REL", 0.5),
-            (1, 3, "REL2", "REL2", 1.5),
+            (0, 2, "Y", {"relationshipType": "Y", "score": 1.0}),
+            (1, 3, "X", {"relationshipType": "X", "weight": 1.5}),
         ]