graphdatascience 1.22__tar.gz → 2.0a1__tar.gz

{graphdatascience-1.22/src/graphdatascience.egg-info → graphdatascience-2.0a1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphdatascience
-Version: 1.22
+Version: 2.0a1
 Summary: A Python client for the Neo4j Graph Data Science (GDS) library
 Author-email: Neo4j <team-gds@neo4j.org>
 License-Expression: Apache-2.0
@@ -27,9 +27,9 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: multimethod<3.0,>=1.0
-Requires-Dist: neo4j<7.0,>=4.4.12
-Requires-Dist: numpy<2.4
-Requires-Dist: pandas<3.0,>=1.0
+Requires-Dist: neo4j<7.0,>=5.26.0
+Requires-Dist: numpy<3.0
+Requires-Dist: pandas<4.0,>=2.0
 Requires-Dist: pyarrow<24.0,>=18.0
 Requires-Dist: textdistance<5.0,>=4.0
 Requires-Dist: tqdm<5.0,>=4.0
@@ -42,7 +42,7 @@ Requires-Dist: networkx<4.0,>=2.0; extra == "networkx"
 Provides-Extra: ogb
 Requires-Dist: ogb<2.0,>=1.0; extra == "ogb"
 Provides-Extra: rust-ext
-Requires-Dist: neo4j-rust-ext<7.0,>=4.4.12; extra == "rust-ext"
+Requires-Dist: neo4j-rust-ext<7.0,>=5.26.0; extra == "rust-ext"
 Dynamic: license-file
 
 # Neo4j Graph Data Science Client

{graphdatascience-1.22 → graphdatascience-2.0a1}/pyproject.toml

@@ -29,9 +29,9 @@ classifiers = [
 requires-python = ">=3.10"
 dependencies = [
     "multimethod >= 1.0, < 3.0",
-    "neo4j >= 4.4.12, < 7.0",
-    "numpy < 2.4",
-    "pandas >= 1.0, < 3.0",
+    "neo4j >= 5.26.0, < 7.0",
+    "numpy < 3.0",
+    "pandas>=2.0,<4.0",
     "pyarrow >= 18.0, < 24.0",
     "textdistance >= 4.0, < 5.0",
     "tqdm >= 4.0, < 5.0",
@@ -50,7 +50,7 @@ BugTracker = "https://github.com/neo4j/graph-data-science-client/issues"
 [project.optional-dependencies]
 networkx = ["networkx >= 2.0, < 4.0"]
 ogb = ["ogb >= 1.0, < 2.0"]
-rust-ext = ["neo4j-rust-ext >= 4.4.12, < 7.0"]
+rust-ext = ["neo4j-rust-ext >= 5.26.0, < 7.0"]
 
 
 [dependency-groups]
@@ -96,22 +96,39 @@ notebook-base = [
     "nbformat==5.10.4",
     "nbclient==0.10.4",
     "ipykernel==7.1.0",
+    "neo4j-viz"
 ]
 notebook-aura-ci = [
     { include-group = "notebook-base" },
-    "python-dotenv==1.2.1"
+    "python-dotenv==1.2.1",
+    "pyspark[sql]"
 ]
 notebook-ci = [
     { include-group = "notebook-base" },
     "scipy == 1.14.0",
     "torch==2.3.0",
-    "torch-scatter==2.1.1",
-    "torch-sparse==0.6.17",
+    "torch-scatter==2.1.2",
+    "torch-sparse==0.6.18",
     "torch-geometric>=2.5.0",
 ]
 
 publish = ["build", "twine"]
 
+# torch-scatter/torch-sparse have no sdist-friendly install: building from source
+# requires a working C++/torch toolchain and fails on many local machines (notably
+# macOS arm64). Instead we pull prebuilt wheels from the PyG wheel index, which ships
+# wheels for both Linux (torch-2.3.0+cpu) and macOS (universal2).
+[tool.uv.sources]
+torch-scatter = [{ index = "pytorch-geometric-cpu" }]
+torch-sparse = [{ index = "pytorch-geometric-cpu" }]
+
+[[tool.uv.index]]
+name = "pytorch-geometric-cpu"
+url = "https://data.pyg.org/whl/torch-2.3.0+cpu.html"
+explicit = true
+format = "flat"
+
+# Fallback only: if uv ever has to build these from source, make torch available.
 [tool.uv.extra-build-dependencies]
 torch-sparse = ["torch"]
 torch-scatter = ["torch"]

graphdatascience-2.0a1/src/graphdatascience/__init__.py

@@ -0,0 +1,11 @@
+from .graph_data_science import GraphDataScience
+from .server_version.server_version import ServerVersion
+from .session import GdsSessions
+from .version import __version__
+
+__all__ = [
+    "__version__",
+    "ServerVersion",
+    "GraphDataScience",
+    "GdsSessions",
+]

{graphdatascience-1.22 → graphdatascience-2.0a1}/src/graphdatascience/arrow_client/arrow_info.py

@@ -4,7 +4,6 @@ from dataclasses import dataclass
 
 from graphdatascience.query_runner.query_runner import QueryRunner
 from graphdatascience.query_runner.query_type import QueryType
-from graphdatascience.server_version.server_version import ServerVersion
 
 
 @dataclass(frozen=True)
@@ -16,12 +15,11 @@ class ArrowInfo:
 
     @staticmethod
     def create(query_runner: QueryRunner) -> ArrowInfo:
-        debugYields = ["listenAddress", "enabled", "running"]
-        if query_runner.server_version() > ServerVersion(2, 6, 0):
-            debugYields.append("versions")
-
         procResult = query_runner.call_procedure(
-            endpoint="gds.debug.arrow", query_type=QueryType.SYSTEM, custom_error=False, yields=debugYields
+            endpoint="gds.debug.arrow",
+            query_type=QueryType.SYSTEM,
+            custom_error=False,
+            yields=["listenAddress", "enabled", "running", "versions"],
         ).iloc[0]
 
         return ArrowInfo(

graphdatascience-2.0a1/src/graphdatascience/arrow_client/arrow_table_utils.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import pandas
+import pyarrow
+from pyarrow import types as pa_types
+
+
+def _downcast_type(data_type: pyarrow.DataType) -> pyarrow.DataType:
+    """Map pyarrow "large" string/binary types back to their 32-bit variants.
+
+    pandas 3.0 introduced a default string dtype that ``pyarrow.Table.from_pandas``
+    maps to ``large_string`` (LargeUTF8), whereas pandas 2.x produced ``string``
+    (UTF8). The GDS server expects UTF8, so we coerce the large variants back.
+    List containers are preserved as-is, only their element type is downcast.
+    """
+    if pa_types.is_large_string(data_type):
+        return pyarrow.string()
+
+    return data_type
+
+
+def table_from_pandas(data: pandas.DataFrame) -> pyarrow.Table:
+    """Convert a DataFrame to a pyarrow Table, coercing large string/binary types to UTF8/binary.
+
+    Use this instead of ``pyarrow.Table.from_pandas`` so the produced schema stays
+    consistent across pandas 2.x and 3.0 (see :func:`_downcast_type`).
+    """
+    table = pyarrow.Table.from_pandas(data)
+
+    new_types = [_downcast_type(field.type) for field in table.schema]
+    if all(new_type == field.type for new_type, field in zip(new_types, table.schema)):
+        return table
+
+    new_schema = pyarrow.schema(
+        [field.with_type(new_type) for field, new_type in zip(table.schema, new_types)],
+        metadata=table.schema.metadata,
+    )
+    return table.cast(new_schema)

{graphdatascience-1.22 → graphdatascience-2.0a1}/src/graphdatascience/arrow_client/authenticated_flight_client.py

@@ -22,7 +22,6 @@ from pyarrow.flight import (
 )
 
 from graphdatascience.arrow_client.arrow_authentication import ArrowAuthentication
-from graphdatascience.arrow_client.arrow_info import ArrowInfo
 from graphdatascience.retry_utils.retry_config import ExponentialWaitConfig, RetryConfigV2, StopConfig
 
 from ..version import __version__
@@ -32,100 +31,66 @@ from .middleware.user_agent_middleware import UserAgentFactory
 
 
 class AuthenticatedArrowClient:
-    @staticmethod
-    def create(
-        arrow_info: ArrowInfo,
-        auth: ArrowAuthentication | None = None,
-        encrypted: bool = False,
-        arrow_client_options: dict[str, Any] | None = None,
-        connection_string_override: str | None = None,
-        retry_config: RetryConfigV2 | None = None,
-        advertised_listen_address: tuple[str, int] | None = None,
-    ) -> AuthenticatedArrowClient:
-        connection_string: str
-        if connection_string_override is not None:
-            connection_string = connection_string_override
-        else:
-            connection_string = arrow_info.listenAddress
-
-        host, port = connection_string.split(":")
-
-        if retry_config is None:
-            retry_config = RetryConfigV2(
-                retryable_exceptions=[
-                    FlightTimedOutError,
-                    FlightUnavailableError,
-                    FlightInternalError,
-                ],
-                stop_config=StopConfig(after_delay=10, after_attempt=5),
-                wait_config=ExponentialWaitConfig(multiplier=1, min=1, max=10),
-            )
-
-        return AuthenticatedArrowClient(
-            host=host,
-            retry_config=retry_config,
-            port=int(port),
-            auth=auth,
-            encrypted=encrypted,
-            arrow_client_options=arrow_client_options,
-            advertised_listen_address=advertised_listen_address,
-        )
-
     def __init__(
         self,
-        host: str,
-        retry_config: RetryConfigV2,
-        port: int = 8491,
+        connection_info: str | tuple[str, int],
         auth: ArrowAuthentication | None = None,
         encrypted: bool = False,
         arrow_client_options: dict[str, Any] | None = None,
         user_agent: str | None = None,
         advertised_listen_address: tuple[str, int] | None = None,
+        retry_config: RetryConfigV2 | None = None,
     ):
         """Creates a new AuthenticatedArrowClient instance.
 
         Parameters
         ----------
-        host: str
-            The host address of the GDS Arrow server
-        port: int
-            The host port of the GDS Arrow server (default is 8491)
-        auth: ArrowAuthentication | None
+        connection_info
+            The host address and port of the GDS Arrow server
+        auth
             An implementation of ArrowAuthentication providing a pair to be used for basic authentication
-        encrypted: bool
+        encrypted
             A flag that indicates whether the connection should be encrypted (default is False)
-        arrow_client_options: dict[str, Any] | None
+        arrow_client_options
             Additional options to be passed to the Arrow Flight client.
-        user_agent: str | None
+        user_agent
             The user agent string to use for the connection. (default is `neo4j-graphdatascience-v[VERSION] pyarrow-v[PYARROW_VERSION])
-        retry_config: RetryConfig | None
+        retry_config
             The retry configuration to use for the Arrow requests send by the client.
-        advertised_listen_address: tuple[str, int] | None
+        advertised_listen_address
             The advertised listen address of the GDS Arrow server. This will be used by remote projection and writeback operations.
         """
+
+        if isinstance(connection_info, str):
+            host, port_str = connection_info.split(":")
+            port = int(port_str)
+        else:
+            host, port = connection_info
+
+        if retry_config is None:
+            retry_config = RetryConfigV2(
+                retryable_exceptions=[
+                    FlightTimedOutError,
+                    FlightUnavailableError,
+                    FlightInternalError,
+                ],
+                stop_config=StopConfig(after_delay=10, after_attempt=5),
+                wait_config=ExponentialWaitConfig(multiplier=1, min=1, max=10),
+            )
+
         self._host = host
-        self._port = port
-        self._auth = None
+        self._port = int(port)
+        self._auth = auth
         self._encrypted = encrypted
         self._arrow_client_options = arrow_client_options
         self._user_agent = user_agent
         self._logger = logging.getLogger("gds_arrow_client")
         self._retry_config = retry_config
         if auth:
-            self._auth = auth
             self._auth_middleware = AuthMiddleware(auth)
         self.advertised_listen_address = advertised_listen_address
-
         self._flight_client: flight.FlightClient = self._instantiate_flight_client()
 
-    def _reconnect(self) -> None:
-        try:
-            self._flight_client.close()
-        except Exception:
-            pass
-
-        self._flight_client = self._instantiate_flight_client()
-
     def connection_info(self) -> ConnectionInfo:
         """
         Returns the host and port of the GDS Arrow server.
@@ -267,6 +232,14 @@ class AuthenticatedArrowClient:
         self.__dict__.update(state)
         self._flight_client = self._instantiate_flight_client()
 
+    def _reconnect(self) -> None:
+        try:
+            self._flight_client.close()
+        except Exception:
+            pass
+
+        self._flight_client = self._instantiate_flight_client()
+
 
 @dataclass
 class ConnectionInfo:

{graphdatascience-1.22 → graphdatascience-2.0a1}/src/graphdatascience/arrow_client/v1/gds_arrow_client.py

@@ -12,11 +12,11 @@ from pyarrow.types import is_dictionary
 from pydantic import BaseModel
 
 from graphdatascience.arrow_client.arrow_endpoint_version import ArrowEndpointVersion
+from graphdatascience.arrow_client.arrow_table_utils import table_from_pandas
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient, ConnectionInfo
 from graphdatascience.arrow_client.v1.data_mapper_utils import deserialize_single
 
 from ...procedure_surface.arrow.error_handler import handle_flight_error
-from ...semantic_version.semantic_version import SemanticVersion
 from ..progress_callback import ProgressCallback
 
 
@@ -487,11 +487,13 @@ class GdsArrowClient:
         batch_size: int,
         progress_callback: ProgressCallback,
     ) -> None:
+        batches: list[RecordBatch]
         match data:
             case pyarrow.Table():
                 batches = data.to_batches(batch_size)
             case pandas.DataFrame():
-                batches = pyarrow.Table.from_pandas(data).to_batches(batch_size)
+                batches = table_from_pandas(data).to_batches(batch_size)
+
             case _:
                 batches = data
 
@@ -562,11 +564,7 @@ class GdsArrowClient:
         except Exception as e:
             handle_flight_error(e)
         arrow_table = self._sanitize_arrow_table(arrow_table)
-        if SemanticVersion.from_string(pandas.__version__) >= SemanticVersion(2, 0, 0):
-            return arrow_table.to_pandas(types_mapper=pandas.ArrowDtype)  # type: ignore
-        else:
-            arrow_table = self._sanitize_arrow_table(arrow_table)
-            return arrow_table.to_pandas()  # type: ignore
+        return arrow_table.to_pandas(types_mapper=pandas.ArrowDtype)  # type: ignore
 
     def __enter__(self) -> GdsArrowClient:
         return self

{graphdatascience-1.22 → graphdatascience-2.0a1}/src/graphdatascience/arrow_client/v2/gds_arrow_client.py

@@ -10,6 +10,7 @@ import pyarrow
 from pyarrow import RecordBatch, flight
 
 from graphdatascience.arrow_client.arrow_endpoint_version import ArrowEndpointVersion
+from graphdatascience.arrow_client.arrow_table_utils import table_from_pandas
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient, ConnectionInfo
 from graphdatascience.query_runner.termination_flag import TerminationFlag
 
@@ -504,14 +505,14 @@ class GdsArrowClient:
         batch_size: int,
     ) -> list[list[pyarrow.RecordBatch]]:
         if isinstance(data, pandas.DataFrame):
-            return [pyarrow.Table.from_pandas(data).to_batches(batch_size)]
+            return [table_from_pandas(data).to_batches(batch_size)]
 
         if isinstance(data, pyarrow.Table):
             return [data.to_batches(batch_size)]
 
         if isinstance(data, list):
             if all(isinstance(entry, pandas.DataFrame) for entry in data):
-                return [pyarrow.Table.from_pandas(entry).to_batches(batch_size) for entry in data]
+                return [table_from_pandas(entry).to_batches(batch_size) for entry in data]
 
             if all(isinstance(entry, pyarrow.RecordBatch) for entry in data):
                 return [data]

{graphdatascience-1.22 → graphdatascience-2.0a1}/src/graphdatascience/arrow_client/v2/job_client.py

@@ -8,7 +8,7 @@ from tenacity import Retrying, retry_if_result
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
 from graphdatascience.arrow_client.v2.api_types import JobIdConfig, JobStatus
 from graphdatascience.arrow_client.v2.data_mapper_utils import deserialize_single
-from graphdatascience.query_runner.progress.progress_bar import TqdmProgressBar
+from graphdatascience.progress.progress_bar import TqdmProgressBar
 from graphdatascience.query_runner.termination_flag import TerminationFlag
 from graphdatascience.retry_utils.retry_utils import job_wait_strategy
 

{graphdatascience-1.22 → graphdatascience-2.0a1}/src/graphdatascience/datasets/graph_constructor_func.py

@@ -2,7 +2,7 @@ from typing import Protocol
 
 from pandas import DataFrame
 
-from graphdatascience.graph.v2 import GraphV2
+from graphdatascience.graph import Graph
 
 
 class GraphConstructorFunc(Protocol):
@@ -13,4 +13,4 @@ class GraphConstructorFunc(Protocol):
         relationships: DataFrame | list[DataFrame] | None = None,
         concurrency: int | None = None,
         undirected_relationship_types: list[str] | None = None,
-    ) -> GraphV2: ...
+    ) -> Graph: ...

{graphdatascience-1.22 → graphdatascience-2.0a1}/src/graphdatascience/datasets/nx_loader.py

@@ -5,14 +5,14 @@ import networkx as nx
 from pandas import DataFrame
 
 from graphdatascience.datasets.graph_constructor_func import GraphConstructorFunc
-from graphdatascience.graph.v2 import GraphV2
+from graphdatascience.graph import Graph
 
 
 class NXLoader:
     def __init__(self, graph_constructor: GraphConstructorFunc) -> None:
         self._construct = graph_constructor
 
-    def load(self, nx_G: nx.Graph, graph_name: str, concurrency: int | None = None) -> GraphV2:
+    def load(self, nx_G: nx.Graph, graph_name: str, concurrency: int | None = None) -> Graph:
         """
         Parameters
         ----------
@@ -25,7 +25,7 @@ class NXLoader:
 
         Returns
         -------
-        GraphV2
+        Graph
             A handle to the graph.
         """
         nodes, rels = self._parse(nx_G)

{graphdatascience-1.22 → graphdatascience-2.0a1}/src/graphdatascience/datasets/ogb_loader.py

@@ -10,7 +10,7 @@ import numpy.typing as npt
 import pandas as pd
 
 from graphdatascience.datasets.graph_constructor_func import GraphConstructorFunc
-from graphdatascience.graph.v2 import GraphV2
+from graphdatascience.graph import Graph
 
 
 class _HomogeneousOGBGraphBase(TypedDict):
@@ -194,7 +194,7 @@ class OGBNLoader(OGBLoader):
         dataset_root_path: str = "dataset",
         graph_name: str | None = None,
         concurrency: int | None = None,
-    ) -> GraphV2:
+    ) -> Graph:
         """
         Parameters
         ----------
@@ -209,7 +209,7 @@ class OGBNLoader(OGBLoader):
 
         Returns
         -------
-        GraphV2
+        Graph
             A handle to the graph.
         """
         try:
@@ -326,7 +326,7 @@ class OGBLLoader(OGBLoader):
         dataset_root_path: str = "dataset",
         graph_name: str | None = None,
         concurrency: int | None = None,
-    ) -> GraphV2:
+    ) -> Graph:
         """
         Parameters
         ----------
@@ -341,7 +341,7 @@ class OGBLLoader(OGBLoader):
 
         Returns
         -------
-        GraphV2
+        Graph
             A handle to the graph.
         """
         try:

{graphdatascience-1.22 → graphdatascience-2.0a1}/src/graphdatascience/datasets/simple_file_loader.py

@@ -1,7 +1,6 @@
 import pathlib
 from typing import NamedTuple
 
-from neo4j import __version__ as neo4j_driver_version
 from pandas import DataFrame, read_parquet
 
 
@@ -11,9 +10,6 @@ class GraphResources(NamedTuple):
 
 
 class SimpleDatasetLoader:
-    def __init__(self) -> None:
-        self._is_neo4j_4_driver = neo4j_driver_version.startswith("4.")
-
     @staticmethod
     def _path(package: str, resource: str) -> pathlib.Path:
         from importlib.resources import files
@@ -25,10 +21,6 @@ class SimpleDatasetLoader:
         file = self._path("graphdatascience.resources.cora", "cora_nodes.parquet.gzip")
         nodes = read_parquet(file)
 
-        if self._is_neo4j_4_driver:
-            # features is read as an ndarray which was not supported in neo4j 4
-            nodes["features"] = nodes["features"].apply(lambda x: x.tolist())
-
         rels = read_parquet(self._path("graphdatascience.resources.cora", "cora_rels.parquet.gzip"))
         return GraphResources([nodes], [rels])
 
@@ -49,9 +41,6 @@ class SimpleDatasetLoader:
         for n in nodes:
             resource = self._path(package, f"imdb_{n}.parquet.gzip")
             df = read_parquet(resource)
-            if self._is_neo4j_4_driver:
-                # features is read as an ndarray which was not supported in neo4j 4
-                df["plot_keywords"] = df["plot_keywords"].apply(lambda x: x.tolist())
             node_dfs.append(df)
 
         rel_dfs = []

{graphdatascience-1.22 → graphdatascience-2.0a1}/src/graphdatascience/error/cypher_warning_handler.py

@@ -2,8 +2,6 @@ import warnings
 from functools import wraps
 from typing import Any, Callable, TypeVar, cast
 
-from ..caller_base import CallerBase
-
 F = TypeVar("F", bound=Callable[..., Any])
 
 
@@ -12,7 +10,7 @@ def filter_id_func_deprecation_warning() -> Callable[[F], F]:
         wraps(func)
 
         @wraps(func)
-        def wrapper(self: CallerBase, *args: Any, **kwargs: Any) -> Any:
+        def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
             # GDS uses the numeric id to resolve the node
             warnings.filterwarnings(
                 "ignore",
@@ -38,6 +36,11 @@ def filter_id_func_deprecation_warning() -> Callable[[F], F]:
                 message=r"warn: feature deprecated without replacement\. id is deprecated and will be removed without a replacement\.",
             )
 
+            warnings.filterwarnings(
+                "ignore",
+                message=r".*is deprecated, use `element_id` instead",
+            )
+
             return func(self, *args, **kwargs)
 
         return cast(F, wrapper)

graphdatascience-2.0a1/src/graphdatascience/error/standalone_session_error.py

@@ -0,0 +1,3 @@
+class NotAvailableInStandaloneSessions(Exception):
+    def __init__(self, subject: str) -> None:
+        super().__init__(self, f"{subject} is not available in standalone sessions")

graphdatascience-2.0a1/src/graphdatascience/graph/__init__.py

@@ -0,0 +1,3 @@
+from graphdatascience.graph.graph_api import Graph
+
+__all__ = ["Graph"]

{graphdatascience-1.22/src/graphdatascience/graph/v2 → graphdatascience-2.0a1/src/graphdatascience/graph}/graph_api.py

@@ -4,11 +4,11 @@ from datetime import datetime
 from types import TracebackType
 from typing import Any, Type
 
-from graphdatascience.graph.v2.graph_backend import GraphBackend
+from graphdatascience.graph.graph_backend import GraphBackend
 from graphdatascience.procedure_surface.api.catalog.graph_info import GraphInfo
 
 
-class GraphV2:
+class Graph:
     """
     A graph object that represents a graph in the graph catalog.
     It can be passed into algorithm endpoints to compute over the corresponding graph.
@@ -19,7 +19,7 @@ class GraphV2:
         self._name = name
         self._backend = backend
 
-    def __enter__(self: GraphV2) -> GraphV2:
+    def __enter__(self: Graph) -> Graph:
         return self
 
     def __exit__(

{graphdatascience-1.22/src/graphdatascience/graph/v2 → graphdatascience-2.0a1/src/graphdatascience/graph}/graph_backend_cypher.py

@@ -1,15 +1,16 @@
 from __future__ import annotations
 
 from graphdatascience.call_parameters import CallParameters
-from graphdatascience.graph.v2 import GraphBackend, GraphV2
+from graphdatascience.graph import Graph
+from graphdatascience.graph.graph_backend import GraphBackend
 from graphdatascience.procedure_surface.api.catalog.graph_info import GraphInfo, GraphInfoWithDegrees
 from graphdatascience.query_runner.query_runner import QueryRunner
 
 
-def get_graph(name: str, query_runner: QueryRunner) -> GraphV2:
+def get_graph(name: str, query_runner: QueryRunner) -> Graph:
     backend = CypherGraphBackend(name, query_runner)
 
-    return GraphV2(name, backend)
+    return Graph(name, backend)
 
 
 class CypherGraphBackend(GraphBackend):