Merge pull request #392 from neo4j/support-gds-2.0

Support GDS 2.0 API

Author: FlorentinD

.github/workflows/gds-integration-tests.yml

@@ -20,6 +20,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest]
+        gds-version: ["1.22", "2.0.0a1"]
     defaults:
       run:
         working-directory: python-wrapper
@@ -33,11 +34,12 @@ jobs:
         with:
           python-version: "3.11"
           enable-cache: true
-      - run: uv sync --group dev --extra pandas --extra neo4j --extra gds
+      - name: Install just
+        uses: extractions/setup-just@v2
 
       - name: Run tests
         env:
           AURA_API_CLIENT_ID: 4V1HYCYEeoU4dSxThKnBeLvE2U4hSphx
           AURA_API_CLIENT_SECRET: ${{ secrets.AURA_API_CLIENT_SECRET }}
           AURA_API_PROJECT_ID: 3f8df5e7-4800-4d4f-ad1d-2d044dfd587c
-        run: uv run pytest tests/ --include-neo4j-and-gds
+        run: just py-ci-test-gds "${{ matrix.gds-version }}"

justfile

@@ -25,6 +25,16 @@ py-test:
     cd python-wrapper && uv sync --all-extras --group dev
     cd python-wrapper && uv run --group dev pytest
 
+# install a specific GDS client version and run the GDS integration tests (used by CI)
+# example: just py-ci-test-gds 2.0.0a1
+py-ci-test-gds gds_version:
+    #!/usr/bin/env bash
+    set -e
+    cd {{py_dir}}
+    uv sync --group dev --extra pandas --extra neo4j --extra gds
+    uv pip install "graphdatascience=={{gds_version}}"
+    uv run pytest tests/ --include-neo4j-and-gds
+
 py-test-gds:
     #!/usr/bin/env bash
     set -e

python-wrapper/pyproject.toml

@@ -43,7 +43,7 @@ requires-python = ">=3.10"
 
 [project.optional-dependencies]
 pandas = ["pandas>=2, <3", "pandas-stubs>=2, <3"]
-gds = ["graphdatascience>=1.22, <2"]
+gds = ["graphdatascience>=1.22, <3"]
 neo4j = ["neo4j"]
 snowflake = ["snowflake-snowpark-python>=1, <2"]
 

python-wrapper/src/neo4j_viz/_gds_compat.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+import importlib
+import re
+from typing import Any
+
+from graphdatascience.version import __version__ as _gds_version
+
+
+def _parse_major(version: str) -> int:
+    match = re.match(r"\s*(\d+)", version)
+    return int(match.group(1)) if match else 0
+
+
+IS_GDS_2: bool = _parse_major(_gds_version) >= 2
+
+if IS_GDS_2:
+    GdsGraph: Any = importlib.import_module("graphdatascience.graph").Graph
+    _GRAPH_TYPES: tuple[type, ...] = (GdsGraph,)
+else:
+    GdsGraph = importlib.import_module("graphdatascience.graph.v2").GraphV2
+    _GRAPH_TYPES = (GdsGraph, importlib.import_module("graphdatascience").Graph)
+
+
+def _check_graph_type(G: Any) -> None:
+    """Raise ``TypeError`` unless ``G`` is a graph object accepted by the installed client."""
+    if not isinstance(G, _GRAPH_TYPES):
+        accepted = " or ".join(t.__name__ for t in _GRAPH_TYPES)
+        raise TypeError(f"`G` must be a GDS graph object ({accepted}), but got {type(G).__name__}")
+
+
+def _catalog(gds: Any) -> Any:
+    """Return the graph catalog endpoints for either client version."""
+    return gds.graph if IS_GDS_2 else gds.v2.graph
+
+
+def _degree_centrality(gds: Any) -> Any:
+    """Return the degree centrality endpoints for either client version."""
+    return gds.degree_centrality if IS_GDS_2 else gds.v2.degree_centrality

python-wrapper/src/neo4j_viz/gds.py

@@ -2,29 +2,29 @@
 
 import warnings
 from itertools import chain
-from typing import Collection, Optional
+from typing import Any, Collection, Optional
 from uuid import uuid4
 
 import pandas as pd
-from graphdatascience import Graph, GraphDataScience
-from graphdatascience.graph.v2 import GraphV2
+from graphdatascience import GraphDataScience
 from graphdatascience.session import AuraGraphDataScience
 
 from neo4j_viz.colors import NEO4J_COLORS_DISCRETE, ColorSpace
 
+from ._gds_compat import IS_GDS_2, GdsGraph, _catalog, _check_graph_type, _degree_centrality
 from .pandas import _from_dfs
 from .visualization_graph import VisualizationGraph
 
 
 def _fetch_node_dfs(
     gds: GraphDataScience | AuraGraphDataScience,
-    G: GraphV2,
+    G: Any,
     node_properties_by_label: dict[str, list[str]],
     node_labels: Collection[str],
     additional_db_node_properties: list[str],
 ) -> dict[str, pd.DataFrame]:
     return {
-        lbl: gds.v2.graph.node_properties.stream(
+        lbl: _catalog(gds).node_properties.stream(
             G,
             node_properties=node_properties_by_label[lbl],
             node_labels=[lbl],
@@ -34,14 +34,14 @@ def _fetch_node_dfs(
     }
 
 
-def _fetch_rel_dfs(gds: GraphDataScience | AuraGraphDataScience, G: GraphV2) -> list[pd.DataFrame]:
+def _fetch_rel_dfs(gds: GraphDataScience | AuraGraphDataScience, G: Any) -> list[pd.DataFrame]:
     rel_props = G.relationship_properties()
 
     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():
-        rel_df = gds.v2.graph.relationships.stream(
+        rel_df = _catalog(gds).relationships.stream(
             G, relationship_types=[rel_type], relationship_properties=list(props)
         )
 
@@ -62,7 +62,7 @@ def _fetch_rel_dfs(gds: GraphDataScience | AuraGraphDataScience, G: GraphV2) ->
 
 def from_gds(
     gds: GraphDataScience | AuraGraphDataScience,
-    G: Graph | GraphV2,
+    G: Any,
     node_properties: Optional[list[str]] = None,
     db_node_properties: Optional[list[str]] = None,
     max_node_count: int = 10_000,
@@ -97,8 +97,9 @@ def from_gds(
     """
     if db_node_properties is None:
         db_node_properties = []
-    if isinstance(G, Graph):
-        G_v2 = gds.v2.graph.get(G.name())
+    _check_graph_type(G)
+    if not IS_GDS_2 and not isinstance(G, GdsGraph):
+        G_v2 = _catalog(gds).get(G.name())
     else:
         G_v2 = G
 
@@ -127,7 +128,7 @@ def from_gds(
         )
         sampling_ratio = float(max_node_count) / node_count
         sample_name = f"neo4j-viz_sample_{uuid4()}"
-        G_fetched, _ = gds.v2.graph.sample.rwr(
+        G_fetched, _ = _catalog(gds).sample.rwr(
             G_v2, sample_name, sampling_ratio=sampling_ratio, node_label_stratification=True
         )
     else:
@@ -139,7 +140,7 @@ def from_gds(
         # as a temporary property to ensure that we have at least one property for each label to fetch
         if sum([len(props) == 0 for props in node_properties_by_label.values()]) > 0:
             property_name = f"neo4j-viz_property_{uuid4()}"
-            gds.v2.degree_centrality.mutate(G_fetched, mutate_property=property_name)
+            _degree_centrality(gds).mutate(G_fetched, mutate_property=property_name)
             for props in node_properties_by_label.values():
                 props.append(property_name)
 
@@ -155,7 +156,7 @@ def from_gds(
         if G_fetched.name() != G.name():
             G_fetched.drop()
         elif property_name is not None:
-            gds.v2.graph.node_properties.drop(G_fetched, node_properties=[property_name])
+            _catalog(gds).node_properties.drop(G_fetched, node_properties=[property_name])
 
     for df in node_dfs.values():
         if property_name is not None and property_name in df.columns:

python-wrapper/tests/neo4j_and_gds/test_gds.py

@@ -1,14 +1,13 @@
 import re
-from contextlib import AbstractContextManager
-from typing import Generator
+from typing import Any, Generator
 
 import pandas as pd
 import pytest
 from graphdatascience import GraphDataScience
-from graphdatascience.graph.v2 import GraphV2
 from graphdatascience.session import AuraGraphDataScience
 
 from neo4j_viz import Node
+from neo4j_viz._gds_compat import _catalog
 from neo4j_viz.gds import from_gds
 
 
@@ -25,11 +24,11 @@ def db_setup(gds: GraphDataScience | AuraGraphDataScience) -> Generator[None, No
     gds.run_cypher("MATCH (n:_CI_A|_CI_B) DETACH DELETE n")
 
 
-def project_graph(gds: GraphDataScience | AuraGraphDataScience) -> AbstractContextManager[GraphV2]:
+def project_graph(gds: GraphDataScience | AuraGraphDataScience) -> Any:
     if isinstance(gds, GraphDataScience):
-        return gds.v2.graph.project("g2", "*", "*")
+        return _catalog(gds).project("g2", "*", "*")
     elif isinstance(gds, AuraGraphDataScience):
-        return gds.v2.graph.project("g2", "MATCH (n)–->(m) RETURN gds.graph.project.remote(n, m)")
+        return _catalog(gds).project("g2", "MATCH (n)–->(m) RETURN gds.graph.project.remote(n, m)")
     raise Exception(f"Unsupported GDS type {type(gds)}")
 
 
@@ -64,7 +63,7 @@ def test_from_gds_integration_all_properties(gds: GraphDataScience | AuraGraphDa
         }
     )
 
-    with gds.v2.graph.construct("flo", nodes, rels) as G:
+    with _catalog(gds).construct("flo", nodes, rels) as G:
         VG = from_gds(gds, G)
 
         assert len(VG.nodes) == 3
@@ -113,7 +112,7 @@ def test_from_gds_integration_all_properties(gds: GraphDataScience | AuraGraphDa
 
 @pytest.mark.requires_neo4j_and_gds
 def test_from_gds_sample(gds: GraphDataScience | AuraGraphDataScience) -> None:
-    with gds.v2.graph.generate("hello", node_count=11_000, average_degree=1) as G:
+    with _catalog(gds).generate("hello", node_count=11_000, average_degree=1) as G:
         with pytest.warns(
             UserWarning,
             match=re.escape(
@@ -164,7 +163,7 @@ def test_from_gds_hetero(gds: GraphDataScience | AuraGraphDataScience) -> None:
         }
     )
 
-    with gds.v2.graph.construct("flo", [A_nodes, B_nodes], [X_rels, Y_rels]) as G:
+    with _catalog(gds).construct("flo", [A_nodes, B_nodes], [X_rels, Y_rels]) as G:
         VG = from_gds(
             gds,
             G,