Allow specifying arrow and writing configuration through proc call

Mats-SX · Mats-SX · commit 85550a9cae64 · 2024-07-02T15:20:26.000+02:00
signatures

this allows for example:

gds.degree.write(G, arrowConfiguration={batchSize: 1})

gds.graph.nodeProperties.write(G, "prop", [], arrowConfiguration={batchSize: 1})
diff --git a/graphdatascience/query_runner/aura_db_query_runner.py b/graphdatascience/query_runner/aura_db_query_runner.py
@@ -104,7 +104,7 @@ def _remote_projection(
         database: Optional[str] = None,
         logging: bool = False,
     ) -> DataFrame:
-        self._inject_connection_parameters(params)
+        self._inject_arrow_config(params["arrow_configuration"])
         return self._db_query_runner.call_procedure(endpoint, params, yields, database, logging, False)
 
     def _remote_write_back(
@@ -119,21 +119,27 @@ def _remote_write_back(
         if params["config"] is None:
             params["config"] = {}
 
+        # we pop these out so that they are not retained for the GDS proc call
+        db_write_config = params["config"].pop("writeConfiguration", {})  # type: ignore
+        db_arrow_config = params["config"].pop("arrowConfiguration", {})  # type: ignore
+        self._inject_write_config(endpoint, params, db_write_config)
+        self._inject_arrow_config(db_arrow_config)
+
         params["config"]["writeToResultStore"] = True  # type: ignore
         gds_write_result = self._gds_query_runner.call_procedure(
             endpoint, params, yields, database, logging, custom_error
         )
 
-        write_params = {
+        db_write_proc_params = {
             "graphName": params["graph_name"],
             "databaseName": self._gds_query_runner.database(),
-            "writeConfiguration": self._extract_write_back_arguments(endpoint, params),
+            "writeConfiguration": db_write_config,
+            "arrowConfiguration": db_arrow_config,
         }
-        self._inject_connection_parameters(write_params)
 
         write_back_start = time.time()
         database_write_result = self._db_query_runner.call_procedure(
-            "gds.arrow.write", CallParameters(write_params), yields, None, False, False
+            "gds.arrow.write", CallParameters(db_write_proc_params), yields, None, False, False
         )
         write_millis = (time.time() - write_back_start) * 1000
         gds_write_result["writeMillis"] = write_millis
@@ -149,22 +155,20 @@ def _remote_write_back(
 
         return gds_write_result
 
-    def _inject_connection_parameters(self, params: Dict[str, Any]) -> None:
+    def _inject_arrow_config(self, params: Dict[str, Any]) -> None:
         host, port = self._gds_arrow_client.connection_info()
         token = self._gds_arrow_client.request_token()
         if token is None:
             token = "IGNORED"
-        params["arrowConfiguration"] = {
-            "host": host,
-            "port": port,
-            "token": token,
-            "encrypted": self._encrypted,
-        }
+
+        params["host"] = host
+        params["port"] = port
+        params["token"] = token
+        params["encrypted"] = self._encrypted
 
     @staticmethod
-    def _extract_write_back_arguments(proc_name: str, params: Dict[str, Any]) -> Dict[str, Any]:
-        config = params.get("config", {})
-        write_config = {}
+    def _inject_write_config(proc_name: str, proc_params: Dict[str, Any], write_config: Dict[str, Any]) -> None:
+        config = proc_params.get("config", {})
 
         if "writeConcurrency" in config:
             write_config["concurrency"] = config["writeConcurrency"]
@@ -188,21 +192,21 @@ def _extract_write_back_arguments(proc_name: str, params: Dict[str, Any]) -> Dic
 
         elif "gds.graph." in proc_name:
             if "gds.graph.nodeProperties.write" == proc_name:
-                properties = params["properties"]
+                properties = proc_params["properties"]
                 write_config["nodeProperties"] = properties if isinstance(properties, list) else [properties]
-                write_config["nodeLabels"] = params["entities"]
+                write_config["nodeLabels"] = proc_params["entities"]
 
             elif "gds.graph.nodeLabel.write" == proc_name:
-                write_config["nodeLabels"] = [params["node_label"]]
+                write_config["nodeLabels"] = [proc_params["node_label"]]
 
             elif "gds.graph.relationshipProperties.write" == proc_name:
-                write_config["relationshipProperties"] = params["relationship_properties"]
-                write_config["relationshipType"] = params["relationship_type"]
+                write_config["relationshipProperties"] = proc_params["relationship_properties"]
+                write_config["relationshipType"] = proc_params["relationship_type"]
 
             elif "gds.graph.relationship.write" == proc_name:
-                if "relationship_property" in params and params["relationship_property"] != "":
-                    write_config["relationshipProperties"] = [params["relationship_property"]]
-                write_config["relationshipType"] = params["relationship_type"]
+                if "relationship_property" in proc_params and proc_params["relationship_property"] != "":
+                    write_config["relationshipProperties"] = [proc_params["relationship_property"]]
+                write_config["relationshipType"] = proc_params["relationship_type"]
 
             else:
                 raise ValueError(f"Unsupported procedure name: {proc_name}")
@@ -215,9 +219,7 @@ def _extract_write_back_arguments(proc_name: str, params: Dict[str, Any]) -> Dic
             else:
                 if "writeProperty" in config:
                     write_config["nodeProperties"] = [config["writeProperty"]]
-                if "nodeLabels" in params:
-                    write_config["nodeLabels"] = params["nodeLabels"]
+                if "nodeLabels" in proc_params:
+                    write_config["nodeLabels"] = proc_params["nodeLabels"]
                 else:
                     write_config["nodeLabels"] = ["*"]
-
-        return write_config
diff --git a/graphdatascience/tests/unit/conftest.py b/graphdatascience/tests/unit/conftest.py
@@ -68,9 +68,13 @@ def encrypted(self) -> bool:
         return False
 
     def last_query(self) -> str:
+        if len(self.queries) == 0:
+            return ""
         return self.queries[-1]
 
-    def last_params(self) -> Dict[str, Any]:
+    def last_params(self) -> dict[str, Any]:
+        if len(self.queries) == 0:
+            return {}
         return self.params[-1]
 
     def set_database(self, database: str) -> None:
diff --git a/graphdatascience/tests/unit/test_aura_db_arrow_query_runner.py b/graphdatascience/tests/unit/test_aura_db_arrow_query_runner.py
@@ -0,0 +1,170 @@
+from typing import Tuple
+
+from pandas import DataFrame
+
+from graphdatascience import ServerVersion
+from graphdatascience.call_parameters import CallParameters
+from graphdatascience.query_runner.aura_db_arrow_query_runner import (
+    AuraDbArrowQueryRunner,
+)
+from graphdatascience.tests.unit.conftest import CollectingQueryRunner
+
+
+class FakeArrowClient:
+
+    def connection_info(self) -> Tuple[str, str]:
+        return "myHost", "1234"
+
+    def request_token(self) -> str:
+        return "myToken"
+
+
+def test_extracts_parameters_projection() -> None:
+    version = ServerVersion(2, 7, 0)
+    db_query_runner = CollectingQueryRunner(version)
+    gds_query_runner = CollectingQueryRunner(version)
+    gds_query_runner.set__mock_result(DataFrame([{"databaseLocation": "remote"}]))
+    qr = AuraDbArrowQueryRunner(gds_query_runner, db_query_runner, FakeArrowClient(), False)  # type: ignore
+
+    qr.call_procedure(
+        endpoint="gds.arrow.project",
+        params=CallParameters(
+            graph_name="g",
+            query="RETURN 1",
+            concurrency=2,
+            undirRels=[],
+            inverseRels=[],
+            arrow_configuration={"batchSize": 100},
+        ),
+    )
+
+    # doesn't run anything on GDS
+    assert gds_query_runner.last_query() == ""
+    assert gds_query_runner.last_params() == {}
+    assert (
+        db_query_runner.last_query()
+        == "CALL gds.arrow.project($graph_name, $query, $concurrency, $undirRels, $inverseRels, $arrow_configuration)"
+    )
+    assert db_query_runner.last_params() == {
+        "graph_name": "g",
+        "query": "RETURN 1",
+        "concurrency": 2,
+        "undirRels": [],
+        "inverseRels": [],
+        "arrow_configuration": {
+            "encrypted": False,
+            "host": "myHost",
+            "port": "1234",
+            "token": "myToken",
+            "batchSize": 100,
+        },
+    }
+
+
+def test_extracts_parameters_algo_write() -> None:
+    version = ServerVersion(2, 7, 0)
+    db_query_runner = CollectingQueryRunner(version)
+    gds_query_runner = CollectingQueryRunner(version)
+    gds_query_runner.set__mock_result(DataFrame([{"databaseLocation": "remote"}]))
+    qr = AuraDbArrowQueryRunner(gds_query_runner, db_query_runner, FakeArrowClient(), False)  # type: ignore
+
+    qr.call_procedure(endpoint="gds.degree.write", params=CallParameters(graph_name="g", config={}))
+
+    assert gds_query_runner.last_query() == "CALL gds.degree.write($graph_name, $config)"
+    assert gds_query_runner.last_params() == {
+        "graph_name": "g",
+        "config": {"writeToResultStore": True},
+    }
+    assert (
+        db_query_runner.last_query()
+        == "CALL gds.arrow.write($graphName, $databaseName, $writeConfiguration, $arrowConfiguration)"
+    )
+    assert db_query_runner.last_params() == {
+        "graphName": "g",
+        "databaseName": "dummy",
+        "writeConfiguration": {"nodeLabels": ["*"]},
+        "arrowConfiguration": {"encrypted": False, "host": "myHost", "port": "1234", "token": "myToken"},
+    }
+
+
+def test_arrow_and_write_configuration() -> None:
+    version = ServerVersion(2, 7, 0)
+    db_query_runner = CollectingQueryRunner(version)
+    gds_query_runner = CollectingQueryRunner(version)
+    gds_query_runner.set__mock_result(DataFrame([{"databaseLocation": "remote"}]))
+    qr = AuraDbArrowQueryRunner(gds_query_runner, db_query_runner, FakeArrowClient(), False)  # type: ignore
+
+    qr.call_procedure(
+        endpoint="gds.degree.write",
+        params=CallParameters(
+            graph_name="g",
+            config={"arrowConfiguration": {"batchSize": 1000}, "writeConfiguration": {"writeMode": "FOOBAR"}},
+        ),
+    )
+
+    assert gds_query_runner.last_query() == "CALL gds.degree.write($graph_name, $config)"
+    assert gds_query_runner.last_params() == {
+        "graph_name": "g",
+        "config": {"writeToResultStore": True},
+    }
+    assert (
+        db_query_runner.last_query()
+        == "CALL gds.arrow.write($graphName, $databaseName, $writeConfiguration, $arrowConfiguration)"
+    )
+    assert db_query_runner.last_params() == {
+        "graphName": "g",
+        "databaseName": "dummy",
+        "writeConfiguration": {"nodeLabels": ["*"], "writeMode": "FOOBAR"},
+        "arrowConfiguration": {
+            "encrypted": False,
+            "host": "myHost",
+            "port": "1234",
+            "token": "myToken",
+            "batchSize": 1000,
+        },
+    }
+
+
+def test_arrow_and_write_configuration_graph_write() -> None:
+    version = ServerVersion(2, 7, 0)
+    db_query_runner = CollectingQueryRunner(version)
+    gds_query_runner = CollectingQueryRunner(version)
+    gds_query_runner.set__mock_result(DataFrame([{"databaseLocation": "remote"}]))
+    qr = AuraDbArrowQueryRunner(gds_query_runner, db_query_runner, FakeArrowClient(), False)  # type: ignore
+
+    qr.call_procedure(
+        endpoint="gds.graph.nodeProperties.write",
+        params=CallParameters(
+            graph_name="g",
+            properties=[],
+            entities=[],
+            config={"arrowConfiguration": {"batchSize": 42}, "writeConfiguration": {"writeMode": "FOOBAR"}},
+        ),
+    )
+
+    assert (
+        gds_query_runner.last_query()
+        == "CALL gds.graph.nodeProperties.write($graph_name, $properties, $entities, $config)"
+    )
+    assert gds_query_runner.last_params() == {
+        "graph_name": "g",
+        "entities": [],
+        "properties": [],
+        "config": {"writeToResultStore": True},
+    }
+    assert (
+        db_query_runner.last_query()
+        == "CALL gds.arrow.write($graphName, $databaseName, $writeConfiguration, $arrowConfiguration)"
+    )
+    assert db_query_runner.last_params() == {
+        "graphName": "g",
+        "databaseName": "dummy",
+        "writeConfiguration": {"nodeLabels": [], "nodeProperties": [], "writeMode": "FOOBAR"},
+        "arrowConfiguration": {
+            "encrypted": False,
+            "host": "myHost",
+            "port": "1234",
+            "token": "myToken",
+            "batchSize": 42,
+        },
+    }
