graphdatascience 1.17a1__tar.gz → 1.18__tar.gz

graphdatascience-1.18/MANIFEST.in

@@ -0,0 +1,3 @@
+prune graphdatascience/resources/cora/serialize_cora.py
+prune graphdatascience/resources/imdb/serialize_imdb.py
+prune graphdatascience/resources/lastfm/serialize_lastfm.py

{graphdatascience-1.17a1/graphdatascience.egg-info → graphdatascience-1.18}/PKG-INFO

@@ -1,22 +1,19 @@
 Metadata-Version: 2.4
 Name: graphdatascience
-Version: 1.17a1
+Version: 1.18
 Summary: A Python client for the Neo4j Graph Data Science (GDS) library
-Home-page: https://neo4j.com/product/graph-data-science/
-Author: Neo4j
-Author-email: team-gds@neo4j.org
-License: Apache License 2.0
+Author-email: Neo4j <team-gds@neo4j.org>
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://neo4j.com/product/graph-data-science/
 Project-URL: Documentation, https://neo4j.com/docs/graph-data-science-client/current/
-Project-URL: Source, https://github.com/neo4j/graph-data-science-client
-Project-URL: Bug Tracker, https://github.com/neo4j/graph-data-science-client/issues
+Project-URL: Repository, https://github.com/neo4j/graph-data-science-client
+Project-URL: BugTracker, https://github.com/neo4j/graph-data-science-client/issues
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
-Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3 :: Only
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
@@ -25,39 +22,27 @@ Classifier: Topic :: Database
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Software Development
 Classifier: Typing :: Typed
-Requires-Python: >=3.9
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: multimethod<3.0,>=1.0
-Requires-Dist: neo4j<6.0,>=4.4.12
-Requires-Dist: numpy<2.3
+Requires-Dist: neo4j<7.0,>=4.4.12
+Requires-Dist: numpy<2.4
 Requires-Dist: pandas<3.0,>=1.0
-Requires-Dist: pyarrow<21.0,>=17.0
+Requires-Dist: pyarrow<23.0,>=17.0
 Requires-Dist: textdistance<5.0,>=4.0
 Requires-Dist: tqdm<5.0,>=4.0
 Requires-Dist: typing-extensions<5.0,>=4.0
 Requires-Dist: requests
 Requires-Dist: tenacity>=9.0
 Requires-Dist: pydantic>=2.11
-Provides-Extra: ogb
-Requires-Dist: ogb<2.0,>=1.0; extra == "ogb"
 Provides-Extra: networkx
 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<6.0,>=4.4.12; extra == "rust-ext"
-Dynamic: author
-Dynamic: author-email
-Dynamic: classifier
-Dynamic: description
-Dynamic: description-content-type
-Dynamic: home-page
-Dynamic: license
+Requires-Dist: neo4j-rust-ext<7.0,>=4.4.12; extra == "rust-ext"
 Dynamic: license-file
-Dynamic: project-url
-Dynamic: provides-extra
-Dynamic: requires-dist
-Dynamic: requires-python
-Dynamic: summary
 
 # Neo4j Graph Data Science Client
 

graphdatascience-1.18/pyproject.toml

@@ -0,0 +1,175 @@
+[build-system]
+requires = ["setuptools>=42", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "graphdatascience"
+description = "A Python client for the Neo4j Graph Data Science (GDS) library"
+dynamic = ["version"]
+authors = [{ name = "Neo4j", email = "team-gds@neo4j.org" }]
+readme = "README.md"
+license = "Apache-2.0"
+classifiers = [
+    "Development Status :: 5 - Production/Stable",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Database",
+    "Topic :: Scientific/Engineering",
+    "Topic :: Software Development",
+    "Typing :: Typed",
+]
+requires-python = ">=3.10"
+dependencies = [
+    "multimethod >= 1.0, < 3.0",
+    "neo4j >= 4.4.12, < 7.0",
+    "numpy < 2.4",
+    "pandas >= 1.0, < 3.0",
+    "pyarrow >= 17.0, < 23.0",
+    "textdistance >= 4.0, < 5.0",
+    "tqdm >= 4.0, < 5.0",
+    "typing-extensions >= 4.0, < 5.0",
+    "requests",
+    "tenacity >= 9.0",
+    "pydantic >= 2.11",
+]
+
+[project.urls]
+Homepage = "https://neo4j.com/product/graph-data-science/"
+Documentation = "https://neo4j.com/docs/graph-data-science-client/current/"
+Repository = "https://github.com/neo4j/graph-data-science-client"
+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"]
+
+
+[dependency-groups]
+# To install all development dependencies,
+# run `pip install --group dev -e .` inside repository root folder.#
+dev = [
+    "tox>4.24",
+    { include-group = "dev-base" },
+    { include-group = "test" },
+    { include-group = "docs-ci" },
+]
+
+
+dev-base = [
+    "ruff == 0.11.7",
+    "mypy == 1.13.0",
+    "nbconvert == 7.16.4",
+    "pandas-stubs == 2.2.3.241009",
+    "tox == 4.30.2",
+    "types-setuptools == 75.8.0.20250110",
+    "enum-tools[sphinx] == 0.12.0",
+    "autodoc_pydantic",
+    "types-requests",
+    "types-tqdm",
+    "types-python-dateutil",
+    "python-dotenv==1.1.0",
+    "types-colorama",
+]
+test = [
+    "pytest == 8.3.3",
+    "requests_mock == 1.11.0",
+    "pytest_mock == 3.15.1",
+    "testcontainers >= 4.0, < 4.13.0",
+    "python-dateutil >= 2.9",
+]
+docs-ci = [
+    "sphinx == 7.3.7",
+    "enum-tools[sphinx] == 0.12.0",
+    "autodoc_pydantic",
+]
+notebook-base = [
+    "nbconvert==7.16.4",
+    "nbformat==5.10.4",
+    "nbclient==0.10.2",
+    "ipykernel==6.29.5",
+]
+notebook-aura-ci = [{ include-group = "notebook-base" }, "python-dotenv==1.1.0"]
+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-geometric>=2.5.0",
+]
+
+publish = ["build", "twine"]
+
+
+[tool.setuptools]
+include-package-data = true
+zip-safe = false
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["graphdatascience*"]
+
+[tool.setuptools.dynamic]
+version = { attr = "graphdatascience.version.__version__" }
+
+[tool.setuptools.package-data]
+graphdatascience = ["py.typed", "resources/**/*.gzip"]
+
+[tool.ruff]
+line-length = 120
+
+include = ["*.py", "*.ipynb"]
+
+exclude = [
+    # default
+
+    ".bzr",
+    ".direnv",
+    ".eggs",
+    ".git",
+    ".git-rewrite",
+    ".hg",
+    ".ipynb_checkpoints",
+    ".mypy_cache",
+    ".nox",
+    ".pants.d",
+    ".pyenv",
+    ".pytest_cache",
+    ".pytype",
+    ".ruff_cache",
+    ".svn",
+    ".tox",
+    ".venv",
+    ".vscode",
+    "__pypackages__",
+    "_build",
+    "buck-out",
+    "build",
+    "dist",
+    "node_modules",
+    "site-packages",
+    "venv",
+
+    # custom
+    "venv*",
+    ".venv*",
+    ".tox",
+    "build",
+]
+
+[tool.ruff.lint]
+select = [
+    "E4",
+    "E7",
+    "E9", # default pycodestyle rules
+    "F",  # flake8
+    "I",  # isort
+]

{graphdatascience-1.17a1 → graphdatascience-1.18/src}/graphdatascience/algo/algo_proc_runner.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from abc import ABC
 from typing import Any
 
@@ -19,7 +21,7 @@ class AlgoProcRunner(IllegalAttrChecker, ABC):
         return self._query_runner.call_procedure(endpoint=self._namespace, params=params, logging=with_logging)
 
     @graph_type_check
-    def estimate(self, G: Graph, **config: Any) -> "Series[Any]":
+    def estimate(self, G: Graph, **config: Any) -> Series[Any]:
         self._namespace += "." + "estimate"
         return self._run_procedure(G, config, with_logging=False).squeeze()  # type: ignore
 
@@ -30,13 +32,13 @@ class StreamModeRunner(AlgoProcRunner):
 
 
 class StandardModeRunner(AlgoProcRunner):
-    def __call__(self, G: Graph, **config: Any) -> "Series[Any]":
+    def __call__(self, G: Graph, **config: Any) -> Series[Any]:
         return self._run_procedure(G, config).squeeze()  # type: ignore
 
 
 class GraphSageRunner(AlgoProcRunner):
     @graph_type_check
-    def __call__(self, G: Graph, **config: Any) -> tuple[GraphSageModel, "Series[Any]"]:
+    def __call__(self, G: Graph, **config: Any) -> tuple[GraphSageModel, Series[Any]]:
         result = self._run_procedure(G, config).squeeze()
         model_name = result["modelInfo"]["modelName"]
 

graphdatascience-1.18/src/graphdatascience/arrow_client/arrow_base_model.py

@@ -0,0 +1,12 @@
+from typing import Any
+
+from pydantic import BaseModel
+from pydantic.alias_generators import to_camel
+
+
+class ArrowBaseModel(BaseModel, alias_generator=to_camel):
+    def dump_camel(self) -> dict[str, Any]:
+        return self.model_dump(by_alias=True)
+
+    def dump_json(self) -> str:
+        return self.model_dump_json(by_alias=True)

graphdatascience-1.18/src/graphdatascience/arrow_client/authenticated_flight_client.py

@@ -0,0 +1,254 @@
+from __future__ import annotations
+
+import json
+import logging
+import platform
+from dataclasses import dataclass
+from types import TracebackType
+from typing import Any, Iterator, Type
+
+import certifi
+from pyarrow import Schema, flight
+from pyarrow import __version__ as arrow_version
+from pyarrow._flight import (
+    Action,
+    ActionType,
+    FlightInternalError,
+    FlightStreamReader,
+    FlightTimedOutError,
+    FlightUnavailableError,
+    Result,
+    Ticket,
+)
+
+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__
+from .middleware.auth_middleware import AuthFactory, AuthMiddleware
+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,
+        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,
+    ):
+        """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
+            An implementation of ArrowAuthentication providing a pair to be used for basic authentication
+        encrypted: bool
+            A flag that indicates whether the connection should be encrypted (default is False)
+        arrow_client_options: dict[str, Any] | None
+            Additional options to be passed to the Arrow Flight client.
+        user_agent: str | None
+            The user agent string to use for the connection. (default is `neo4j-graphdatascience-v[VERSION] pyarrow-v[PYARROW_VERSION])
+        retry_config: RetryConfig | None
+            The retry configuration to use for the Arrow requests send by the client.
+        advertised_listen_address: tuple[str, int] | None
+            The advertised listen address of the GDS Arrow server. This will be used by remote projection and writeback operations.
+        """
+        self._host = host
+        self._port = port
+        self._auth = None
+        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 connection_info(self) -> ConnectionInfo:
+        """
+        Returns the host and port of the GDS Arrow server.
+
+        Returns
+        -------
+        ConnectionInfo
+            the host and port of the GDS Arrow server
+        """
+        return ConnectionInfo(self._host, self._port, self._encrypted)
+
+    def advertised_connection_info(self) -> ConnectionInfo:
+        """
+        Returns the advertised host and port of the GDS Arrow server.
+
+        Returns
+        -------
+        ConnectionInfo
+            the host and port of the GDS Arrow server
+        """
+        if self.advertised_listen_address is None:
+            return self.connection_info()
+
+        h, p = self.advertised_listen_address
+        return ConnectionInfo(h, p, self._encrypted)
+
+    def request_token(self) -> str | None:
+        """
+        Requests a token from the server and returns it.
+
+        Returns
+        -------
+        str | None
+            a token from the server and returns it.
+        """
+
+        @self._retry_config.decorator(operation_name="Request token", logger=self._logger)
+        def auth_with_retry() -> None:
+            client = self._flight_client
+            if self._auth:
+                auth_pair = self._auth.auth_pair()
+                client.authenticate_basic_token(auth_pair[0], auth_pair[1])
+
+        if self._auth:
+            auth_with_retry()
+            return self._auth_middleware.token()
+        else:
+            return "IGNORED"
+
+    def get_stream(self, ticket: Ticket) -> FlightStreamReader:
+        return self._flight_client.do_get(ticket)
+
+    def do_action(self, endpoint: str, payload: bytes | dict[str, Any]) -> Iterator[Result]:
+        payload_bytes = payload if isinstance(payload, bytes) else json.dumps(payload).encode("utf-8")
+
+        return self._flight_client.do_action(Action(endpoint, payload_bytes))  # type: ignore
+
+    def do_action_with_retry(self, endpoint: str, payload: bytes | dict[str, Any]) -> list[Result]:
+        @self._retry_config.decorator(operation_name="Send action", logger=self._logger)
+        def run_with_retry() -> list[Result]:
+            # the Flight response error code is only checked on iterator consumption
+            # we eagerly collect iterator here to trigger retry in case of an error
+            return list(self.do_action(endpoint, payload))
+
+        return run_with_retry()
+
+    def list_actions(self) -> set[ActionType]:
+        return self._flight_client.list_actions()  # type: ignore
+
+    def do_put_with_retry(
+        self, descriptor: flight.FlightDescriptor, schema: Schema
+    ) -> tuple[flight.FlightStreamWriter, flight.FlightMetadataReader]:
+        @self._retry_config.decorator(operation_name="Do put", logger=self._logger)
+        def run_with_retry() -> tuple[flight.FlightStreamWriter, flight.FlightMetadataReader]:
+            return self._flight_client.do_put(descriptor, schema)  # type: ignore
+
+        return run_with_retry()
+
+    def __enter__(self) -> AuthenticatedArrowClient:
+        return self
+
+    def __exit__(
+        self,
+        exception_type: Type[BaseException] | None,
+        exception_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self.close()
+
+    def close(self) -> None:
+        if self._flight_client:
+            self._flight_client.close()
+
+    def _instantiate_flight_client(self) -> flight.FlightClient:
+        location = (
+            flight.Location.for_grpc_tls(self._host, self._port)
+            if self._encrypted
+            else flight.Location.for_grpc_tcp(self._host, self._port)
+        )
+        client_options: dict[str, Any] = (self._arrow_client_options or {}).copy()
+
+        # We need to specify the system root certificates on Windows
+        if platform.system() == "Windows":
+            if not client_options["tls_root_certs"]:
+                client_options["tls_root_certs"] = certifi.contents()
+
+        if self._auth:
+            user_agent = f"neo4j-graphdatascience-v{__version__} pyarrow-v{arrow_version}"
+            if self._user_agent:
+                user_agent = self._user_agent
+
+            client_options["middleware"] = [
+                AuthFactory(self._auth_middleware),
+                UserAgentFactory(useragent=user_agent),
+            ]
+
+        return flight.FlightClient(location, **client_options)
+
+    def __getstate__(self) -> dict[str, Any]:
+        state = self.__dict__.copy()
+        # Remove the FlightClient as it isn't serializable
+        if "_flight_client" in state:
+            del state["_flight_client"]
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        self._flight_client = self._instantiate_flight_client()
+
+
+@dataclass
+class ConnectionInfo:
+    host: str
+    port: int
+    encrypted: bool

graphdatascience-1.18/src/graphdatascience/arrow_client/middleware/auth_middleware.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import base64
+import time
+from typing import Any
+
+from pyarrow._flight import ClientMiddleware, ClientMiddlewareFactory
+
+from graphdatascience.arrow_client.arrow_authentication import ArrowAuthentication
+
+
+class AuthFactory(ClientMiddlewareFactory):  # type: ignore
+    def __init__(self, middleware: AuthMiddleware, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._middleware = middleware
+
+    def start_call(self, info: Any) -> AuthMiddleware:
+        return self._middleware
+
+
+class AuthMiddleware(ClientMiddleware):  # type: ignore
+    def __init__(self, auth: ArrowAuthentication, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._auth = auth
+        self._token: str | None = None
+        self._token_timestamp = 0
+
+    def token(self) -> str | None:
+        # check whether the token is older than 10 minutes. If so, reset it.
+        if self._token and int(time.time()) - self._token_timestamp > 600:
+            self._token = None
+
+        return self._token
+
+    def _set_token(self, token: str) -> None:
+        self._token = token
+        self._token_timestamp = int(time.time())
+
+    def received_headers(self, headers: dict[str, Any]) -> None:
+        auth_header = headers.get("authorization", None)
+        if not auth_header:
+            return
+
+        # the result is always a list
+        header_value = auth_header[0]
+
+        if not isinstance(header_value, str):
+            raise ValueError(f"Incompatible header value received from server: `{header_value}`")
+
+        auth_type, token = header_value.split(" ", 1)
+        if auth_type == "Bearer":
+            self._set_token(token)
+
+    def sending_headers(self) -> dict[str, str]:
+        token = self.token()
+        if token is not None:
+            return {"authorization": "Bearer " + token}
+
+        auth_pair = self._auth.auth_pair()
+        auth_token = f"{auth_pair[0]}:{auth_pair[1]}"
+        auth_token = "Basic " + base64.b64encode(auth_token.encode("utf-8")).decode("ASCII")
+        # There seems to be a bug, `authorization` must be lower key
+        return {"authorization": auth_token}

graphdatascience-1.18/src/graphdatascience/arrow_client/middleware/user_agent_middleware.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pyarrow._flight import ClientMiddleware, ClientMiddlewareFactory
+
+
+class UserAgentFactory(ClientMiddlewareFactory):  # type: ignore
+    def __init__(self, useragent: str, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._middleware = UserAgentMiddleware(useragent)
+
+    def start_call(self, info: Any) -> ClientMiddleware:
+        return self._middleware
+
+
+class UserAgentMiddleware(ClientMiddleware):  # type: ignore
+    def __init__(self, useragent: str, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._useragent = useragent
+
+    def sending_headers(self) -> dict[str, str]:
+        return {"x-gds-user-agent": self._useragent}
+
+    def received_headers(self, headers: dict[str, Any]) -> None:
+        pass

graphdatascience-1.18/src/graphdatascience/arrow_client/v1/data_mapper_utils.py

@@ -0,0 +1,19 @@
+import json
+from typing import Any
+
+from pyarrow._flight import Result
+
+
+def deserialize_single(input_stream: list[Result]) -> dict[str, Any]:
+    rows = deserialize(input_stream)
+    if len(rows) != 1:
+        raise ValueError(f"Expected exactly one result, got {len(rows)}")
+
+    return rows[0]
+
+
+def deserialize(input_stream: list[Result]) -> list[dict[str, Any]]:
+    def deserialize_row(row: Result):  # type:ignore
+        return json.loads(row.body.to_pybytes().decode())
+
+    return [deserialize_row(row) for row in input_stream]