Add option for for streaming with Arrow with gds.<algo>.stream

Author: adamnsch

graphdatascience/algo/algo_proc_runner.py

@@ -1,13 +1,21 @@
 from abc import ABC
 from typing import Any, Dict, Tuple
+from uuid import uuid4
+from warnings import warn
 
 from pandas import DataFrame, Series
 
+from ..call_parameters import CallParameters
 from ..error.illegal_attr_checker import IllegalAttrChecker
+from ..graph.graph_entity_ops_runner import (
+    GraphElementPropertyRunner,
+    GraphNodePropertiesRunner,
+    GraphRelationshipsRunner,
+)
 from ..graph.graph_object import Graph
 from ..graph.graph_type_check import graph_type_check
 from ..model.graphsage_model import GraphSageModel
-from graphdatascience.call_parameters import CallParameters
+from ..query_runner.arrow_query_runner import ArrowQueryRunner
 
 
 class AlgoProcRunner(IllegalAttrChecker, ABC):
@@ -24,8 +32,71 @@ def estimate(self, G: Graph, **config: Any) -> "Series[Any]":
 
 
 class StreamModeRunner(AlgoProcRunner):
-    def __call__(self, G: Graph, **config: Any) -> DataFrame:
-        return self._run_procedure(G, config)
+    def __call__(self, G: Graph, stream_with_arrow: bool = False, **config: Any) -> DataFrame:
+        if stream_with_arrow:
+            if not isinstance(self._query_runner, ArrowQueryRunner):
+                raise ValueError("The `stream_with_arrow` option requires GDS EE with the Arrow server enabled")
+
+            try:
+                return self._stream_with_arrow(G, config)
+            except Exception as e:
+                warn(
+                    "Falling back to streaming with Neo4j driver, "
+                    f"since failed to stream with Arrow with reason: {str(e)}"
+                )
+
+                del config["mutateProperty"]
+                del config["mutateRelationshipType"]
+
+                self._namespace = self._namespace.replace("mutate", "stream")
+                return self._run_procedure(G, config)
+        else:
+            return self._run_procedure(G, config)
+
+    def _stream_with_arrow(self, G: Graph, config: Dict[str, Any]) -> DataFrame:
+        self._namespace = self._namespace.replace("stream", "mutate")
+
+        mutate_property = str(uuid4())
+        config["mutateProperty"] = mutate_property
+        try:
+            self._run_procedure(G, config)
+
+            elem_prop_runner = GraphElementPropertyRunner(
+                self._query_runner, "gds.graph.nodeProperty", self._server_version
+            )
+            if "concurrency" in config:
+                result = elem_prop_runner.stream(G, mutate_property, concurrency=config["concurrency"])
+            else:
+                result = elem_prop_runner.stream(G, mutate_property)
+
+            node_prop_runner = GraphNodePropertiesRunner(
+                self._query_runner, "gds.graph.nodeProperties", self._server_version
+            )
+            node_prop_runner.drop(G, [mutate_property])
+
+            return result
+        except Exception as e:
+            if "No value specified for the mandatory configuration parameter `mutateRelationshipType`" not in str(e):
+                raise e
+
+            mutate_relationship_type = str(uuid4())
+            config["mutateRelationshipType"] = mutate_relationship_type
+            self._run_procedure(G, config)
+
+            elem_prop_runner = GraphElementPropertyRunner(
+                self._query_runner, "gds.graph.relationshipProperty", self._server_version
+            )
+            if "concurrency" in config:
+                result = elem_prop_runner.stream(
+                    G, mutate_property, [mutate_relationship_type], concurrency=config["concurrency"]
+                )
+            else:
+                result = elem_prop_runner.stream(G, mutate_property, [mutate_relationship_type])
+
+            rel_runner = GraphRelationshipsRunner(self._query_runner, "gds.graph.relationships", self._server_version)
+            rel_runner.drop(G, mutate_relationship_type)
+
+            return result.drop("relationshipType", axis=1)
 
 
 class StandardModeRunner(AlgoProcRunner):

graphdatascience/tests/integration/test_simple_algo.py

@@ -113,3 +113,21 @@ def test_fastRP_write_estimate(gds: GraphDataScience) -> None:
     result = gds.fastRP.write.estimate(G, writeProperty="embedding", embeddingDimension=4, randomSeed=42)
 
     assert result["requiredMemory"]
+
+
+def test_nodeSimilarity_stream_with_arrow(gds: GraphDataScience) -> None:
+    G, _ = gds.graph.project(GRAPH_NAME, "*", "*")
+
+    result = gds.nodeSimilarity.stream(G, similarityCutoff=0, stream_with_arrow=True)
+
+    assert len(result) == 2
+    assert result["propertyValue"][0] == 0.5
+
+
+def test_fastRP_stream_with_arrow(gds: GraphDataScience) -> None:
+    G, _ = gds.graph.project(GRAPH_NAME, "*", "*")
+
+    result = gds.fastRP.stream(G, stream_with_arrow=True, embeddingDimension=32)
+
+    assert len(result) == 3
+    assert len(result["propertyValue"][0]) == 32