From e31a3113751395b5579caad0b73143c331e2253c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Florentin=20D=C3=B6rre?= Date: Fri, 26 Jun 2026 14:07:44 +0200 Subject: [PATCH 1/2] Support GDS 2.0 API ref GDSA-1212 --- .github/workflows/gds-integration-tests.yml | 6 +- justfile | 10 ++++ python-wrapper/pyproject.toml | 2 +- python-wrapper/src/neo4j_viz/_gds_compat.py | 57 +++++++++++++++++++ python-wrapper/src/neo4j_viz/gds.py | 27 ++++----- .../tests/neo4j_and_gds/test_gds.py | 17 +++--- 6 files changed, 94 insertions(+), 25 deletions(-) create mode 100644 python-wrapper/src/neo4j_viz/_gds_compat.py diff --git a/.github/workflows/gds-integration-tests.yml b/.github/workflows/gds-integration-tests.yml index 53a94b8..69a40f3 100644 --- a/.github/workflows/gds-integration-tests.yml +++ b/.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 }}" diff --git a/justfile b/justfile index 21446a7..7a8a995 100644 --- a/justfile +++ b/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 diff --git a/python-wrapper/pyproject.toml b/python-wrapper/pyproject.toml index 986119d..330bbd7 100644 --- a/python-wrapper/pyproject.toml +++ b/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"] diff --git a/python-wrapper/src/neo4j_viz/_gds_compat.py b/python-wrapper/src/neo4j_viz/_gds_compat.py new file mode 100644 index 0000000..b643a2b --- /dev/null +++ b/python-wrapper/src/neo4j_viz/_gds_compat.py @@ -0,0 +1,57 @@ +"""Compatibility shims for the Neo4j GDS Python client. + +Supports both the GDS 1.22 transitional API (graph operations exposed under the +``gds.v2.*`` namespace, graph objects of type ``GraphV2``) and the GDS 2.0 API +(the v2 endpoints became the default, so they live directly under ``gds.*`` and +the graph class was renamed from ``GraphV2`` to ``Graph``). + +The two APIs are selected by the installed client's major version, exposed via +``graphdatascience.version.__version__``: major ``>= 2`` means the v2 endpoints +are the default. +""" + +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 + + +# In GDS 2.0 the (formerly ``v2``) endpoints became the default: graph operations moved +# from ``gds.v2.*`` to ``gds.*`` and ``GraphV2`` was renamed to ``Graph``. +IS_GDS_2: bool = _parse_major(_gds_version) >= 2 + +# The native graph class for the installed client version. Resolved dynamically (typed +# ``Any``) because its import path differs between versions and only one path exists at a time. +# ``_GRAPH_TYPES`` is the full set of graph objects accepted as input to ``from_gds``: on the +# 1.22 transitional client we also accept the legacy v1 ``Graph`` (which is then converted). +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 diff --git a/python-wrapper/src/neo4j_viz/gds.py b/python-wrapper/src/neo4j_viz/gds.py index 78e37ba..67a1a42 100644 --- a/python-wrapper/src/neo4j_viz/gds.py +++ b/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: diff --git a/python-wrapper/tests/neo4j_and_gds/test_gds.py b/python-wrapper/tests/neo4j_and_gds/test_gds.py index 46d9994..dfbdeab 100644 --- a/python-wrapper/tests/neo4j_and_gds/test_gds.py +++ b/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, From bc16da9365c10379536ec178ccf4fbeb49ede050 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Florentin=20D=C3=B6rre?= Date: Fri, 26 Jun 2026 14:28:49 +0200 Subject: [PATCH 2/2] Reduce comments --- python-wrapper/src/neo4j_viz/_gds_compat.py | 18 ------------------ 1 file changed, 18 deletions(-) diff --git a/python-wrapper/src/neo4j_viz/_gds_compat.py b/python-wrapper/src/neo4j_viz/_gds_compat.py index b643a2b..d09ccb8 100644 --- a/python-wrapper/src/neo4j_viz/_gds_compat.py +++ b/python-wrapper/src/neo4j_viz/_gds_compat.py @@ -1,15 +1,3 @@ -"""Compatibility shims for the Neo4j GDS Python client. - -Supports both the GDS 1.22 transitional API (graph operations exposed under the -``gds.v2.*`` namespace, graph objects of type ``GraphV2``) and the GDS 2.0 API -(the v2 endpoints became the default, so they live directly under ``gds.*`` and -the graph class was renamed from ``GraphV2`` to ``Graph``). - -The two APIs are selected by the installed client's major version, exposed via -``graphdatascience.version.__version__``: major ``>= 2`` means the v2 endpoints -are the default. -""" - from __future__ import annotations import importlib @@ -24,14 +12,8 @@ def _parse_major(version: str) -> int: return int(match.group(1)) if match else 0 -# In GDS 2.0 the (formerly ``v2``) endpoints became the default: graph operations moved -# from ``gds.v2.*`` to ``gds.*`` and ``GraphV2`` was renamed to ``Graph``. IS_GDS_2: bool = _parse_major(_gds_version) >= 2 -# The native graph class for the installed client version. Resolved dynamically (typed -# ``Any``) because its import path differs between versions and only one path exists at a time. -# ``_GRAPH_TYPES`` is the full set of graph objects accepted as input to ``from_gds``: on the -# 1.22 transitional client we also accept the legacy v1 ``Graph`` (which is then converted). if IS_GDS_2: GdsGraph: Any = importlib.import_module("graphdatascience.graph").Graph _GRAPH_TYPES: tuple[type, ...] = (GdsGraph,)