Merge pull request #583 from s1ck/arrow-version-support

Arrow Server version support

Author: s1ck

graphdatascience/error/endpoint_suggester.py

@@ -9,7 +9,7 @@ def generate_suggestive_error_message(requested_endpoint: str, all_endpoints: Li
     MIN_SIMILARITY_FOR_SUGGESTION = 0.9
 
     closest_endpoint = None
-    curr_max_similarity = 0
+    curr_max_similarity = 0.0
     for ep in all_endpoints:
         similarity = textdistance.jaro_winkler(requested_endpoint, ep)
         if similarity >= MIN_SIMILARITY_FOR_SUGGESTION:

graphdatascience/query_runner/arrow_endpoint_version.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import List
+
+
+class ArrowEndpointVersion(Enum):
+    ALPHA = ""
+    V1 = "v1/"
+
+    def version(self) -> str:
+        return self._name_.lower()
+
+    def prefix(self) -> str:
+        return self._value_
+
+    @staticmethod
+    def from_arrow_info(supported_arrow_versions: List[str]) -> ArrowEndpointVersion:
+        # Fallback for pre 2.6.0 servers that do not support versions
+        if len(supported_arrow_versions) == 0:
+            return ArrowEndpointVersion.ALPHA
+
+        # If the server supports versioned endpoints, we try v1 first
+        if ArrowEndpointVersion.V1.version() in supported_arrow_versions:
+            return ArrowEndpointVersion.V1
+
+        if ArrowEndpointVersion.ALPHA.version() in supported_arrow_versions:
+            return ArrowEndpointVersion.ALPHA
+
+        raise UnsupportedArrowEndpointVersion(supported_arrow_versions)
+
+
+class UnsupportedArrowEndpointVersion(Exception):
+    def __init__(self, server_version: List[str]) -> None:
+        super().__init__(self, f"Unsupported Arrow endpoint versions: {server_version}")

graphdatascience/query_runner/arrow_graph_constructor.py

@@ -11,6 +11,7 @@
 from pyarrow import Table
 from tqdm.auto import tqdm
 
+from .arrow_endpoint_version import ArrowEndpointVersion
 from .graph_constructor import GraphConstructor
 
 
@@ -21,13 +22,15 @@ def __init__(
         graph_name: str,
         flight_client: flight.FlightClient,
         concurrency: int,
+        arrow_endpoint_version: ArrowEndpointVersion,
         undirected_relationship_types: Optional[List[str]],
         chunk_size: int = 10_000,
     ):
         self._database = database
         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
         )
@@ -81,6 +84,7 @@ def _partition_dfs(self, dfs: List[DataFrame]) -> List[DataFrame]:
         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
@@ -93,6 +97,7 @@ def _send_df(self, df: DataFrame, entity_type: str, pbar: tqdm) -> 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"))
@@ -117,3 +122,17 @@ def _send_dfs(self, dfs: List[DataFrame], entity_type: str) -> None:
                 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/query_runner/arrow_query_runner.py

@@ -1,17 +1,20 @@
+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, Series
+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 .graph_constructor import GraphConstructor
 from .query_runner import QueryRunner
@@ -28,18 +31,14 @@ def create(
         encrypted: bool = False,
         disable_server_verification: bool = False,
         tls_root_certs: Optional[bytes] = None,
-    ) -> "QueryRunner":
+    ) -> 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()
+        listen_address: str = arrow_info.get("advertisedListenAddress", arrow_info["listenAddress"])
+        arrow_endpoint_version = ArrowEndpointVersion.from_arrow_info(arrow_info.get("versions", []))
 
-        yield_fields = (
-            ["running", "listenAddress"]
-            if server_version >= ServerVersion(2, 2, 1)
-            else ["running", "advertisedListenAddress"]
-        )
-        arrow_info: "Series[Any]" = fallback_query_runner.call_procedure(
-            endpoint="gds.debug.arrow", yields=yield_fields, custom_error=False
-        ).squeeze()
-        listen_address: str = arrow_info.get("advertisedListenAddress", arrow_info["listenAddress"])  # type: ignore
         if arrow_info["running"]:
             return ArrowQueryRunner(
                 listen_address,
@@ -49,6 +48,7 @@ def create(
                 encrypted,
                 disable_server_verification,
                 tls_root_certs,
+                arrow_endpoint_version,
             )
         else:
             return fallback_query_runner
@@ -62,9 +62,11 @@ def __init__(
         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._server_version = server_version
+        self._arrow_endpoint_version = arrow_endpoint_version
 
         host, port_string = uri.split(":")
 
@@ -272,8 +274,15 @@ def _run_arrow_property_get(self, graph_name: str, procedure_name: str, configur
             "procedure_name": procedure_name,
             "configuration": configuration,
         }
-        ticket = flight.Ticket(json.dumps(payload).encode("utf-8"))
 
+        if self._arrow_endpoint_version == ArrowEndpointVersion.V1:
+            payload = {
+                "name": "GET_MESSAGE",
+                "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()
 
@@ -302,7 +311,12 @@ def create_graph_constructor(
             )
 
         return ArrowGraphConstructor(
-            database, graph_name, self._flight_client, concurrency, undirected_relationship_types
+            database,
+            graph_name,
+            self._flight_client,
+            concurrency,
+            self._arrow_endpoint_version,
+            undirected_relationship_types,
         )
 
     def _sanitize_arrow_table(self, arrow_table: Table) -> Table:

graphdatascience/tests/unit/test_arrow_endpoint_version.py

@@ -0,0 +1,38 @@
+from typing import List, Tuple
+
+import pytest
+
+from graphdatascience.query_runner.arrow_endpoint_version import (
+    ArrowEndpointVersion,
+    UnsupportedArrowEndpointVersion,
+)
+
+
+@pytest.mark.parametrize(
+    "arrow_versions",
+    [
+        (ArrowEndpointVersion.ALPHA, []),
+        (ArrowEndpointVersion.ALPHA, ["alpha"]),
+        (ArrowEndpointVersion.V1, ["v1"]),
+        (ArrowEndpointVersion.V1, ["alpha", "v1"]),
+        (ArrowEndpointVersion.V1, ["v1", "v2"]),
+        (ArrowEndpointVersion.ALPHA, ["alpha"]),
+        (ArrowEndpointVersion.ALPHA, ["alpha", "v2"]),
+    ],
+)
+def test_from_arrow_info_multiple_versions(arrow_versions: Tuple[ArrowEndpointVersion, List[str]]) -> None:
+    actual = ArrowEndpointVersion.from_arrow_info(arrow_versions[1])
+    assert actual == arrow_versions[0]
+
+
+@pytest.mark.parametrize("arrow_version", ["v2", "unsupported"])
+def test_from_arrow_info_fails(arrow_version: str) -> None:
+    with pytest.raises(UnsupportedArrowEndpointVersion) as e:
+        ArrowEndpointVersion.from_arrow_info([arrow_version])
+    assert "Unsupported" in str(e.value)
+    assert arrow_version in str(e.value)
+
+
+def test_prefix() -> None:
+    assert ArrowEndpointVersion.ALPHA.prefix() == ""
+    assert ArrowEndpointVersion.V1.prefix() == "v1/"