graphdatascience 1.11a1__tar.gz → 1.11a3__tar.gz

{graphdatascience-1.11a1/graphdatascience.egg-info → graphdatascience-1.11a3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphdatascience
-Version: 1.11a1
+Version: 1.11a3
 Summary: A Python client for the Neo4j Graph Data Science (GDS) library
 Home-page: https://neo4j.com/product/graph-data-science/
 Author: Neo4j
@@ -30,8 +30,9 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: multimethod<2.0,>=1.0
 Requires-Dist: neo4j<6.0,>=4.4.2
+Requires-Dist: numpy<2.0
 Requires-Dist: pandas<3.0,>=1.0
-Requires-Dist: pyarrow<16.0,>=11.0
+Requires-Dist: pyarrow<17.0,>=14.0.1
 Requires-Dist: textdistance<5.0,>=4.0
 Requires-Dist: tqdm<5.0,>=4.0
 Requires-Dist: typing-extensions<5.0,>=4.0

{graphdatascience-1.11a1 → graphdatascience-1.11a3}/graphdatascience/graph/base_graph_proc_runner.py

@@ -15,12 +15,12 @@ from ..error.uncallable_namespace import UncallableNamespace
 from ..server_version.compatible_with import compatible_with
 from ..server_version.server_version import ServerVersion
 from .graph_entity_ops_runner import (
-    GraphElementPropertyRunner,
     GraphLabelRunner,
     GraphNodePropertiesRunner,
     GraphNodePropertyRunner,
     GraphPropertyRunner,
     GraphRelationshipPropertiesRunner,
+    GraphRelationshipPropertyRunner,
     GraphRelationshipRunner,
     GraphRelationshipsRunner,
 )
@@ -390,9 +390,9 @@ class BaseGraphProcRunner(UncallableNamespace, IllegalAttrChecker):
         return GraphNodePropertiesRunner(self._query_runner, self._namespace, self._server_version)
 
     @property
-    def relationshipProperty(self) -> GraphElementPropertyRunner:
+    def relationshipProperty(self) -> GraphRelationshipPropertyRunner:
         self._namespace += ".relationshipProperty"
-        return GraphElementPropertyRunner(self._query_runner, self._namespace, self._server_version)
+        return GraphRelationshipPropertyRunner(self._query_runner, self._namespace, self._server_version)
 
     @property
     def relationshipProperties(self) -> GraphRelationshipPropertiesRunner:

{graphdatascience-1.11a1 → graphdatascience-1.11a3}/graphdatascience/graph/graph_entity_ops_runner.py

@@ -70,13 +70,6 @@ class GraphEntityOpsBaseRunner(UncallableNamespace, IllegalAttrChecker):
         )
 
 
-class GraphElementPropertyRunner(GraphEntityOpsBaseRunner):
-    @compatible_with("stream", min_inclusive=ServerVersion(2, 2, 0))
-    def stream(self, G: Graph, node_properties: str, node_labels: Strings = ["*"], **config: Any) -> DataFrame:
-        self._namespace += ".stream"
-        return self._handle_properties(G, node_properties, node_labels, config)
-
-
 class GraphNodePropertyRunner(GraphEntityOpsBaseRunner):
     @compatible_with("stream", min_inclusive=ServerVersion(2, 2, 0))
     @filter_id_func_deprecation_warning()
@@ -177,7 +170,7 @@ class GraphNodePropertiesRunner(GraphEntityOpsBaseRunner):
         return reduce(add_property, db_node_properties, query_prefix)
 
     @compatible_with("write", min_inclusive=ServerVersion(2, 2, 0))
-    def write(self, G: Graph, node_properties: List[str], node_labels: Strings = ["*"], **config: Any) -> "Series[Any]":
+    def write(self, G: Graph, node_properties: Strings, node_labels: Strings = ["*"], **config: Any) -> "Series[Any]":
         self._namespace += ".write"
         return self._handle_properties(G, node_properties, node_labels, config).squeeze()  # type: ignore
 
@@ -197,6 +190,16 @@ class GraphNodePropertiesRunner(GraphEntityOpsBaseRunner):
         ).squeeze()
 
 
+class GraphRelationshipPropertyRunner(GraphEntityOpsBaseRunner):
+    @compatible_with("stream", min_inclusive=ServerVersion(2, 2, 0))
+    def stream(
+        self, G: Graph, relationship_property: str, relationship_types: Strings = ["*"], **config: Any
+    ) -> DataFrame:
+        self._namespace += ".stream"
+        relationship_types = [relationship_types] if isinstance(relationship_types, str) else relationship_types
+        return self._handle_properties(G, relationship_property, relationship_types, config)
+
+
 class GraphRelationshipPropertiesRunner(GraphEntityOpsBaseRunner):
     @compatible_with("stream", min_inclusive=ServerVersion(2, 2, 0))
     def stream(
@@ -209,6 +212,8 @@ class GraphRelationshipPropertiesRunner(GraphEntityOpsBaseRunner):
     ) -> DataFrame:
         self._namespace += ".stream"
 
+        relationship_types = [relationship_types] if isinstance(relationship_types, str) else relationship_types
+
         result = self._handle_properties(G, relationship_properties, relationship_types, config)
 
         # new format was requested, but the query was run via Cypher

{graphdatascience-1.11a1 → graphdatascience-1.11a3}/graphdatascience/graph/graph_remote_proc_runner.py

@@ -5,5 +5,4 @@ from graphdatascience.graph.graph_remote_project_runner import GraphProjectRemot
 class GraphRemoteProcRunner(BaseGraphProcRunner):
     @property
     def project(self) -> GraphProjectRemoteRunner:
-        self._namespace += ".project.remoteDb"
         return GraphProjectRemoteRunner(self._query_runner, self._namespace, self._server_version)

graphdatascience-1.11a3/graphdatascience/graph/graph_remote_project_runner.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from typing import List, Optional
+
+from ..error.illegal_attr_checker import IllegalAttrChecker
+from ..query_runner.aura_db_query_runner import AuraDbQueryRunner
+from ..server_version.compatible_with import compatible_with
+from .graph_object import Graph
+from graphdatascience.call_parameters import CallParameters
+from graphdatascience.graph.graph_create_result import GraphCreateResult
+from graphdatascience.server_version.server_version import ServerVersion
+
+
+class GraphProjectRemoteRunner(IllegalAttrChecker):
+    @compatible_with("project", min_inclusive=ServerVersion(2, 7, 0))
+    def __call__(
+        self,
+        graph_name: str,
+        query: str,
+        concurrency: int = 4,
+        undirected_relationship_types: Optional[List[str]] = None,
+        inverse_indexed_relationship_types: Optional[List[str]] = None,
+        batch_size: Optional[int] = None,
+    ) -> GraphCreateResult:
+        if inverse_indexed_relationship_types is None:
+            inverse_indexed_relationship_types = []
+        if undirected_relationship_types is None:
+            undirected_relationship_types = []
+
+        arrow_configuration = {}
+        if batch_size is not None:
+            arrow_configuration["batchSize"] = batch_size
+
+        params = CallParameters(
+            graph_name=graph_name,
+            query=query,
+            concurrency=concurrency,
+            undirected_relationship_types=undirected_relationship_types,
+            inverse_indexed_relationship_types=inverse_indexed_relationship_types,
+            arrow_configuration=arrow_configuration,
+        )
+
+        result = self._query_runner.call_procedure(
+            endpoint=AuraDbQueryRunner.GDS_REMOTE_PROJECTION_PROC_NAME,
+            params=params,
+        ).squeeze()
+        return GraphCreateResult(Graph(graph_name, self._query_runner, self._server_version), result)

{graphdatascience-1.11a1 → graphdatascience-1.11a3}/graphdatascience/query_runner/arrow_graph_constructor.py

@@ -1,19 +1,17 @@
 from __future__ import annotations
 
 import concurrent
-import json
 import math
 import warnings
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Dict, List, NoReturn, Optional
 
 import numpy
-import pyarrow.flight as flight
 from pandas import DataFrame
 from pyarrow import Table
 from tqdm.auto import tqdm
 
-from .arrow_endpoint_version import ArrowEndpointVersion
+from .gds_arrow_client import GdsArrowClient
 from .graph_constructor import GraphConstructor
 
 
@@ -22,9 +20,8 @@ class ArrowGraphConstructor(GraphConstructor):
         self,
         database: str,
         graph_name: str,
-        flight_client: flight.FlightClient,
+        flight_client: GdsArrowClient,
         concurrency: int,
-        arrow_endpoint_version: ArrowEndpointVersion,
         undirected_relationship_types: Optional[List[str]],
         chunk_size: int = 10_000,
     ):
@@ -32,7 +29,6 @@ class ArrowGraphConstructor(GraphConstructor):
         self._concurrency = concurrency
         self._graph_name = graph_name
         self._client = flight_client
-        self._arrow_endpoint_version = arrow_endpoint_version
         self._undirected_relationship_types = (
             [] if undirected_relationship_types is None else undirected_relationship_types
         )
@@ -49,20 +45,20 @@ class ArrowGraphConstructor(GraphConstructor):
             if self._undirected_relationship_types:
                 config["undirected_relationship_types"] = self._undirected_relationship_types
 
-            self._send_action(
+            self._client.send_action(
                 "CREATE_GRAPH",
                 config,
             )
 
             self._send_dfs(node_dfs, "node")
 
-            self._send_action("NODE_LOAD_DONE", {"name": self._graph_name})
+            self._client.send_action("NODE_LOAD_DONE", {"name": self._graph_name})
 
             self._send_dfs(relationship_dfs, "relationship")
 
-            self._send_action("RELATIONSHIP_LOAD_DONE", {"name": self._graph_name})
+            self._client.send_action("RELATIONSHIP_LOAD_DONE", {"name": self._graph_name})
         except (Exception, KeyboardInterrupt) as e:
-            self._send_action("ABORT", {"name": self._graph_name})
+            self._client.send_action("ABORT", {"name": self._graph_name})
 
             raise e
 
@@ -85,25 +81,12 @@ class ArrowGraphConstructor(GraphConstructor):
 
         return partitioned_dfs
 
-    def _send_action(self, action_type: str, meta_data: Dict[str, Any]) -> None:
-        action_type = self._versioned_action_type(action_type)
-        result = self._client.do_action(flight.Action(action_type, json.dumps(meta_data).encode("utf-8")))
-
-        # Consume result fully to sanity check and avoid cancelled streams
-        collected_result = list(result)
-        assert len(collected_result) == 1
-
-        json.loads(collected_result[0].body.to_pybytes().decode())
-
     def _send_df(self, df: DataFrame, entity_type: str, pbar: tqdm[NoReturn]) -> None:
         table = Table.from_pandas(df)
         batches = table.to_batches(self._chunk_size)
         flight_descriptor = {"name": self._graph_name, "entity_type": entity_type}
-        flight_descriptor = self._versioned_flight_desriptor(flight_descriptor)
 
-        # Write schema
-        upload_descriptor = flight.FlightDescriptor.for_command(json.dumps(flight_descriptor).encode("utf-8"))
-        writer, _ = self._client.do_put(upload_descriptor, table.schema)
+        writer, _ = self._client.start_put(flight_descriptor, table.schema)
 
         with writer:
             # Write table in chunks
@@ -126,17 +109,3 @@ class ArrowGraphConstructor(GraphConstructor):
                 if not future.exception():
                     continue
                 raise future.exception()  # type: ignore
-
-    def _versioned_action_type(self, action_type: str) -> str:
-        return self._arrow_endpoint_version.prefix() + action_type
-
-    def _versioned_flight_desriptor(self, flight_descriptor: Dict[str, Any]) -> Dict[str, Any]:
-        return (
-            flight_descriptor
-            if self._arrow_endpoint_version == ArrowEndpointVersion.ALPHA
-            else {
-                "name": "PUT_MESSAGE",
-                "version": ArrowEndpointVersion.V1.version(),
-                "body": flight_descriptor,
-            }
-        )

{graphdatascience-1.11a1 → graphdatascience-1.11a3}/graphdatascience/query_runner/arrow_query_runner.py

@@ -1,21 +1,14 @@
 from __future__ import annotations
 
-import base64
-import json
-import time
 import warnings
 from typing import Any, Dict, List, Optional, Tuple
 
-import pyarrow.flight as flight
 from pandas import DataFrame
-from pyarrow import ChunkedArray, Table, chunked_array
-from pyarrow.flight import ClientMiddleware, ClientMiddlewareFactory
-from pyarrow.types import is_dictionary  # type: ignore
 
 from ..call_parameters import CallParameters
 from ..server_version.server_version import ServerVersion
-from .arrow_endpoint_version import ArrowEndpointVersion
 from .arrow_graph_constructor import ArrowGraphConstructor
+from .gds_arrow_client import GdsArrowClient
 from .graph_constructor import GraphConstructor
 from .query_runner import QueryRunner
 from graphdatascience.server_version.compatible_with import (
@@ -33,61 +26,29 @@ class ArrowQueryRunner(QueryRunner):
         tls_root_certs: Optional[bytes] = None,
         connection_string_override: Optional[str] = None,
     ) -> QueryRunner:
-        arrow_info = (
-            fallback_query_runner.call_procedure(endpoint="gds.debug.arrow", custom_error=False).squeeze().to_dict()
-        )
-        server_version = fallback_query_runner.server_version()
-        connection_string: str
-        if connection_string_override is not None:
-            connection_string = connection_string_override
-        else:
-            connection_string = arrow_info.get("advertisedListenAddress", arrow_info["listenAddress"])
-        arrow_endpoint_version = ArrowEndpointVersion.from_arrow_info(arrow_info.get("versions", []))
-
-        if arrow_info["running"]:
-            return ArrowQueryRunner(
-                connection_string,
-                fallback_query_runner,
-                server_version,
-                auth,
-                encrypted,
-                disable_server_verification,
-                tls_root_certs,
-                arrow_endpoint_version,
-            )
-        else:
+        if not GdsArrowClient.is_arrow_enabled(fallback_query_runner):
             return fallback_query_runner
 
+        gds_arrow_client = GdsArrowClient.create(
+            fallback_query_runner,
+            auth,
+            encrypted,
+            disable_server_verification,
+            tls_root_certs,
+            connection_string_override,
+        )
+
+        return ArrowQueryRunner(gds_arrow_client, fallback_query_runner, fallback_query_runner.server_version())
+
     def __init__(
         self,
-        uri: str,
+        gds_arrow_client: GdsArrowClient,
         fallback_query_runner: QueryRunner,
         server_version: ServerVersion,
-        auth: Optional[Tuple[str, str]] = None,
-        encrypted: bool = False,
-        disable_server_verification: bool = False,
-        tls_root_certs: Optional[bytes] = None,
-        arrow_endpoint_version: ArrowEndpointVersion = ArrowEndpointVersion.ALPHA,
     ):
         self._fallback_query_runner = fallback_query_runner
+        self._gds_arrow_client = gds_arrow_client
         self._server_version = server_version
-        self._arrow_endpoint_version = arrow_endpoint_version
-
-        host, port_string = uri.split(":")
-
-        location = (
-            flight.Location.for_grpc_tls(host, int(port_string))
-            if encrypted
-            else flight.Location.for_grpc_tcp(host, int(port_string))
-        )
-
-        client_options: Dict[str, Any] = {"disable_server_verification": disable_server_verification}
-        if auth:
-            client_options["middleware"] = [AuthFactory(auth)]
-        if tls_root_certs:
-            client_options["tls_root_certs"] = tls_root_certs
-
-        self._flight_client = flight.FlightClient(location, **client_options)
 
     def warn_about_deprecation(self, old_endpoint: str, new_endpoint: str) -> None:
         warnings.warn(
@@ -140,7 +101,7 @@ class ArrowQueryRunner(QueryRunner):
                         old_endpoint="gds.graph.streamNodeProperty", new_endpoint="gds.graph.nodeProperty.stream"
                     )
 
-            return self._run_arrow_property_get(graph_name, endpoint, config)
+            return self._gds_arrow_client.get_property(self.database(), graph_name, endpoint, config)
         elif (
             old_endpoint := ("gds.graph.streamNodeProperties" == endpoint)
         ) or "gds.graph.nodeProperties.stream" == endpoint:
@@ -159,7 +120,8 @@ class ArrowQueryRunner(QueryRunner):
                     self.warn_about_deprecation(
                         old_endpoint="gds.graph.streamNodeProperties", new_endpoint="gds.graph.nodeProperties.stream"
                     )
-            return self._run_arrow_property_get(
+            return self._gds_arrow_client.get_property(
+                self.database(),
                 graph_name,
                 endpoint,
                 config,
@@ -180,7 +142,8 @@ class ArrowQueryRunner(QueryRunner):
                         old_endpoint="gds.graph.streamRelationshipProperty",
                         new_endpoint="gds.graph.relationshipProperty.stream",
                     )
-            return self._run_arrow_property_get(
+            return self._gds_arrow_client.get_property(
+                self.database(),
                 graph_name,
                 endpoint,
                 {"relationship_property": property_name, "relationship_types": relationship_types},
@@ -202,7 +165,8 @@ class ArrowQueryRunner(QueryRunner):
                         new_endpoint="gds.graph.relationshipProperties.stream",
                     )
 
-            return self._run_arrow_property_get(
+            return self._gds_arrow_client.get_property(
+                self.database(),
                 graph_name,
                 endpoint,
                 {"relationship_properties": property_names, "relationship_types": relationship_types},
@@ -229,7 +193,9 @@ class ArrowQueryRunner(QueryRunner):
                             new_endpoint="gds.graph.relationships.stream",
                         )
 
-            return self._run_arrow_property_get(graph_name, endpoint, {"relationship_types": relationship_types})
+            return self._gds_arrow_client.get_property(
+                self.database(), graph_name, endpoint, {"relationship_types": relationship_types}
+            )
 
         return self._fallback_query_runner.call_procedure(endpoint, params, yields, database, logging, custom_error)
 
@@ -259,52 +225,11 @@ class ArrowQueryRunner(QueryRunner):
 
     def close(self) -> None:
         self._fallback_query_runner.close()
-        # PyArrow 7 did not expose a close method yet
-        if hasattr(self._flight_client, "close"):
-            self._flight_client.close()
+        self._gds_arrow_client.close()
 
     def fallback_query_runner(self) -> QueryRunner:
         return self._fallback_query_runner
 
-    def _run_arrow_property_get(self, graph_name: str, procedure_name: str, configuration: Dict[str, Any]) -> DataFrame:
-        if not self.database():
-            raise ValueError(
-                "For this call you must have explicitly specified a valid Neo4j database to execute on, "
-                "using `GraphDataScience.set_database`."
-            )
-
-        payload = {
-            "database_name": self.database(),
-            "graph_name": graph_name,
-            "procedure_name": procedure_name,
-            "configuration": configuration,
-        }
-
-        if self._arrow_endpoint_version == ArrowEndpointVersion.V1:
-            payload = {
-                "name": "GET_COMMAND",
-                "version": ArrowEndpointVersion.V1.version(),
-                "body": payload,
-            }
-
-        ticket = flight.Ticket(json.dumps(payload).encode("utf-8"))
-        get = self._flight_client.do_get(ticket)
-        arrow_table = get.read_all()
-
-        if configuration.get("list_node_labels", False):
-            # GDS 2.5 had an inconsistent naming of the node labels column
-            new_colum_names = ["nodeLabels" if i == "labels" else i for i in arrow_table.column_names]
-            arrow_table = arrow_table.rename_columns(new_colum_names)
-
-        # Pandas 2.2.0 deprecated an API used by ArrowTable.to_pandas() (< pyarrow 15.0)
-        warnings.filterwarnings(
-            "ignore",
-            category=DeprecationWarning,
-            message=r"Passing a BlockManager to DataFrame is deprecated",
-        )
-
-        return self._sanitize_arrow_table(arrow_table).to_pandas()  # type: ignore
-
     def create_graph_constructor(
         self, graph_name: str, concurrency: int, undirected_relationship_types: Optional[List[str]]
     ) -> GraphConstructor:
@@ -318,86 +243,7 @@ class ArrowQueryRunner(QueryRunner):
         return ArrowGraphConstructor(
             database,
             graph_name,
-            self._flight_client,
+            self._gds_arrow_client,
             concurrency,
-            self._arrow_endpoint_version,
             undirected_relationship_types,
         )
-
-    def _sanitize_arrow_table(self, arrow_table: Table) -> Table:
-        # empty columns cannot be used to build a chunked_array in pyarrow
-        if len(arrow_table) == 0:
-            return arrow_table
-
-        dict_encoded_fields = [
-            (idx, field) for idx, field in enumerate(arrow_table.schema) if is_dictionary(field.type)
-        ]
-        for idx, field in dict_encoded_fields:
-            try:
-                field.type.to_pandas_dtype()
-            except NotImplementedError:
-                # we need to decode the dictionary column before transforming to pandas
-                if isinstance(arrow_table[field.name], ChunkedArray):
-                    decoded_col = chunked_array([chunk.dictionary_decode() for chunk in arrow_table[field.name].chunks])
-                else:
-                    decoded_col = arrow_table[field.name].dictionary_decode()
-                arrow_table = arrow_table.set_column(idx, field.name, decoded_col)
-        return arrow_table
-
-
-class AuthFactory(ClientMiddlewareFactory):  # type: ignore
-    def __init__(self, auth: Tuple[str, str], *args: Any, **kwargs: Any) -> None:
-        super().__init__(*args, **kwargs)
-        self._auth = auth
-        self._token: Optional[str] = None
-        self._token_timestamp = 0
-
-    def start_call(self, info: Any) -> "AuthMiddleware":
-        return AuthMiddleware(self)
-
-    def token(self) -> Optional[str]:
-        # 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())
-
-    @property
-    def auth(self) -> Tuple[str, str]:
-        return self._auth
-
-
-class AuthMiddleware(ClientMiddleware):  # type: ignore
-    def __init__(self, factory: AuthFactory, *args: Any, **kwargs: Any) -> None:
-        super().__init__(*args, **kwargs)
-        self._factory = factory
-
-    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._factory.set_token(token)
-
-    def sending_headers(self) -> Dict[str, str]:
-        token = self._factory.token()
-        if not token:
-            username, password = self._factory.auth
-            auth_token = f"{username}:{password}"
-            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}
-        else:
-            return {"authorization": "Bearer " + token}