Merge pull request #138 from soerenreichardt/stream-topology-endpoint

Stream topology endpoint

Author: adamnsch

graphdatascience/graph/graph_entity_ops_runner.py

@@ -151,6 +151,15 @@ def drop(
 
         return self._query_runner.run_query(query, params).squeeze()  # type: ignore
 
+    @compatible_with("stream", min_inclusive=ServerVersion(2, 2, 0))
+    def stream(self, G: Graph, relationship_types: List[str] = ["*"], **config: Any) -> DataFrame:
+        self._namespace += ".stream"
+        query = f"CALL {self._namespace}($graph_name, $relationship_types, $config)"
+
+        params = {"graph_name": G.name(), "relationship_types": relationship_types, "config": config}
+
+        return self._query_runner.run_query(query, params)
+
 
 class GraphPropertyRunner(CallerBase, UncallableNamespace, IllegalAttrChecker):
     @compatible_with("stream", min_inclusive=ServerVersion(2, 2, 0))

graphdatascience/query_runner/arrow_query_runner.py

@@ -10,6 +10,9 @@
 from .arrow_graph_constructor import ArrowGraphConstructor
 from .graph_constructor import GraphConstructor
 from .query_runner import QueryRunner
+from graphdatascience.server_version.compatible_with import (
+    IncompatibleServerVersionError,
+)
 from graphdatascience.server_version.server_version import ServerVersion
 
 
@@ -107,6 +110,19 @@ def run_query(self, query: str, params: Optional[Dict[str, Any]] = None) -> Data
                 endpoint,
                 {"relationship_properties": property_names, "relationship_types": relationship_types},
             )
+        elif "gds.beta.graph.relationships.stream" in query:
+            graph_name = params["graph_name"]
+            relationship_types = params["relationship_types"]
+
+            if self._server_version < new_endpoint_server_version:
+                raise IncompatibleServerVersionError(
+                    f"The call gds.beta.graph.relationships.stream with parameters {params} via Arrow requires GDS "
+                    f"server version >= 2.2.0. The current version is {self._server_version}"
+                )
+            else:
+                endpoint = "gds.beta.graph.relationships.stream"
+
+            return self._run_arrow_property_get(graph_name, endpoint, {"relationship_types": relationship_types})
 
         return self._fallback_query_runner.run_query(query, params)
 

graphdatascience/tests/integration/test_coverage.py

@@ -8,7 +8,6 @@
 IGNORED_ENDPOINTS = {
     "gds.alpha.graph.removeGraphProperty",  # Exists but undocumented for GDS 2.1
     "gds.alpha.graph.streamGraphProperty",  # Exists but undocumented for GDS 2.1
-    "gds.beta.graph.relationships.stream",  # FIXME: Add support
     "gds.alpha.pipeline.linkPrediction.addMLP",
     "gds.alpha.pipeline.linkPrediction.addRandomForest",
     "gds.beta.pipeline.linkPrediction.addFeature",

graphdatascience/tests/integration/test_graph_ops.py

@@ -549,6 +549,32 @@ def test_graph_relationshipProperties_stream_without_arrow_separate_property_col
     assert {e for e in result["relY"]} == {5, 6, 7}
 
 
+@pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 2, 0))
+def test_graph_relationships_stream_without_arrow(gds_without_arrow: GraphDataScience) -> None:
+    G, _ = gds_without_arrow.graph.project(GRAPH_NAME, "*", "REL")
+
+    result = gds_without_arrow.beta.graph.relationships.stream(G, ["REL"])
+
+    expected = gds_without_arrow.run_cypher("MATCH (n)-[REL]->(m) RETURN id(n) AS src_id, id(m) AS trg_id")
+
+    assert list(result.keys()) == ["sourceNodeId", "targetNodeId", "relationshipType"]
+    assert {e for e in result["sourceNodeId"]} == {i for i in expected["src_id"]}
+    assert {e for e in result["targetNodeId"]} == {i for i in expected["trg_id"]}
+
+
+@pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 2, 0))
+def test_graph_relationships_stream_with_arrow(gds: GraphDataScience) -> None:
+    G, _ = gds.graph.project(GRAPH_NAME, "*", "REL")
+
+    result = gds.beta.graph.relationships.stream(G, ["REL"])
+
+    expected = gds.run_cypher("MATCH (n)-[REL]->(m) RETURN id(n) AS src_id, id(m) AS trg_id")
+
+    assert list(result.keys()) == ["sourceNodeId", "targetNodeId", "relationshipType"]
+    assert {e for e in result["sourceNodeId"]} == {i for i in expected["src_id"]}
+    assert {e for e in result["targetNodeId"]} == {i for i in expected["trg_id"]}
+
+
 def test_graph_writeNodeProperties(gds: GraphDataScience) -> None:
     G, _ = gds.graph.project(GRAPH_NAME, "*", "*")
 

graphdatascience/tests/unit/test_graph_ops.py

@@ -533,6 +533,15 @@ def test_graph_property_drop(runner: CollectingQueryRunner, gds: GraphDataScienc
     assert runner.last_params() == {"graph_name": "g", "graph_property": "prop", "config": {}}
 
 
+@pytest.mark.parametrize("server_version", [ServerVersion(2, 2, 0)])
+def test_graph_relationships_stream(runner: CollectingQueryRunner, gds: GraphDataScience) -> None:
+    G, _ = gds.graph.project("g", "*", "*")
+
+    gds.beta.graph.relationships.stream(G, ["REL_A"])
+    assert runner.last_query() == "CALL gds.beta.graph.relationships.stream($graph_name, $relationship_types, $config)"
+    assert runner.last_params() == {"graph_name": "g", "relationship_types": ["REL_A"], "config": {}}
+
+
 def test_graph_generate(runner: CollectingQueryRunner, gds: GraphDataScience) -> None:
     gds.beta.graph.generate("g", 1337, 42, orientation="NATURAL")