Merge pull request #878 from neo4j/arrow-flakyness

Fixes to enable diagnose of flaky `graph.construct` test using Arrow

Author: Mats-SX

graphdatascience/query_runner/arrow_graph_constructor.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import concurrent
+import logging
 import math
 import warnings
 from concurrent.futures import ThreadPoolExecutor
@@ -33,6 +34,7 @@ def __init__(
         )
         self._chunk_size = chunk_size
         self._min_batch_size = chunk_size * 10
+        self._logger = logging.getLogger()
 
     def run(self, node_dfs: list[DataFrame], relationship_dfs: list[DataFrame]) -> None:
         try:
@@ -60,8 +62,11 @@ def run(self, node_dfs: list[DataFrame], relationship_dfs: list[DataFrame]) -> N
 
             self._client.relationship_load_done(self._graph_name)
         except (Exception, KeyboardInterrupt) as e:
-            self._client.abort(self._graph_name)
-
+            try:
+                self._client.abort(self._graph_name)
+            except Exception as abort_exception:
+                if "No arrow process" not in str(abort_exception):
+                    self._logger.warning(f"error aborting graph creation: {abort_exception}")
             raise e
 
     def _partition_dfs(self, dfs: list[DataFrame]) -> list[DataFrame]:

graphdatascience/query_runner/gds_arrow_client.py

@@ -640,7 +640,10 @@ def _send_action(self, action_type: str, meta_data: dict[str, Any]) -> dict[str,
         )
         def send_with_retry() -> dict[str, Any]:
             try:
-                result = client.do_action(flight.Action(action_type, json.dumps(meta_data).encode("utf-8")))
+                result = client.do_action(
+                    action=flight.Action(action_type, json.dumps(meta_data).encode("utf-8")),
+                    options=flight.FlightCallOptions(timeout=20.0),
+                )
 
                 # Consume result fully to sanity check and avoid cancelled streams
                 collected_result = list(result)

graphdatascience/tests/integration/test_graph_construct.py

@@ -145,10 +145,9 @@ def test_imdb_graph_with_arrow(gds: GraphDataScience) -> None:
         G.drop()
 
 
-@pytest.mark.filterwarnings("ignore: GDS Enterprise users can use Apache Arrow")
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 3, 0))
-def test_lastfm_graph_without_arrow(gds_without_arrow: GraphDataScience) -> None:
-    G = gds_without_arrow.graph.load_lastfm()
+def test_lastfm_graph_with_arrow(gds: GraphDataScience) -> None:
+    G = gds.graph.load_lastfm()
 
     try:
         assert G.node_count() == 19914
@@ -162,14 +161,14 @@ def test_lastfm_graph_without_arrow(gds_without_arrow: GraphDataScience) -> None
         assert set(G.relationship_properties("TAGGED")) == {"day", "month", "year", "tagID", "timestamp"}
         assert set(G.relationship_properties("LISTEN_TO")) == {"weight"}
         assert set(G.relationship_properties("IS_FRIEND")) == set()
-
     finally:
         G.drop()
 
 
+@pytest.mark.filterwarnings("ignore: GDS Enterprise users can use Apache Arrow")
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 3, 0))
-def test_lastfm_graph_with_arrow(gds: GraphDataScience) -> None:
-    G = gds.graph.load_lastfm()
+def test_lastfm_graph_without_arrow(gds_without_arrow: GraphDataScience) -> None:
+    G = gds_without_arrow.graph.load_lastfm()
 
     try:
         assert G.node_count() == 19914
@@ -183,6 +182,7 @@ def test_lastfm_graph_with_arrow(gds: GraphDataScience) -> None:
         assert set(G.relationship_properties("TAGGED")) == {"day", "month", "year", "tagID", "timestamp"}
         assert set(G.relationship_properties("LISTEN_TO")) == {"weight"}
         assert set(G.relationship_properties("IS_FRIEND")) == set()
+
     finally:
         G.drop()