Set mutateMillis

Co-authored-by: Max Kießling <max.kiessling@neo4j.com>

Author: FlorentinD

graphdatascience/arrow_client/v2/api_types.py

@@ -12,5 +12,6 @@ class JobStatus(ArrowBaseModel):
 
 
 class MutateResult(ArrowBaseModel):
+    mutate_millis: int
     node_properties_written: int
     relationships_written: int

graphdatascience/arrow_client/v2/mutation_client.py

@@ -1,4 +1,6 @@
 import json
+import math
+import time
 
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
 from graphdatascience.arrow_client.v2.api_types import MutateResult
@@ -12,5 +14,7 @@ class MutationClient:
     def mutate_node_property(client: AuthenticatedArrowClient, job_id: str, mutate_property: str) -> MutateResult:
         mutate_config = {"jobId": job_id, "mutateProperty": mutate_property}
         encoded_config = json.dumps(mutate_config).encode("utf-8")
+        start_time = time.time()
         mutate_arrow_res = client.do_action_with_retry(MutationClient.MUTATE_ENDPOINT, encoded_config)
-        return MutateResult(**deserialize_single(mutate_arrow_res))
+        mutate_millis = math.ceil((time.time() - start_time) * 1000)
+        return MutateResult(mutateMillis=mutate_millis, **deserialize_single(mutate_arrow_res))

graphdatascience/procedure_surface/arrow/k1coloring_arrow_endpoints.py

@@ -46,12 +46,12 @@ def mutate(
             username=username,
         )
 
-        computation_result, node_properties_written = self._node_property_endpoints.run_job_and_mutate(
+        computation_result, mutate_result = self._node_property_endpoints.run_job_and_mutate(
             "v2/community.k1coloring", G, config, mutate_property
         )
 
-        computation_result["nodeCount"] = node_properties_written
-        computation_result["mutateMillis"] = 0
+        computation_result["nodeCount"] = mutate_result.node_properties_written
+        computation_result["mutateMillis"] = mutate_result.mutate_millis
 
         return K1ColoringMutateResult(**computation_result)
 

graphdatascience/procedure_surface/arrow/kcore_arrow_endpoints.py

@@ -36,12 +36,12 @@ def mutate(
             sudo=sudo,
         )
 
-        computation_result, node_properties_written = self._node_property_endpoints.run_job_and_mutate(
+        computation_result, mutate_result = self._node_property_endpoints.run_job_and_mutate(
             "v2/community.kcore", G, config, mutate_property
         )
 
-        computation_result["nodePropertiesWritten"] = node_properties_written
-        computation_result["mutateMillis"] = 0
+        computation_result["nodePropertiesWritten"] = mutate_result.node_properties_written
+        computation_result["mutateMillis"] = mutate_result.mutate_millis
 
         return KCoreMutateResult(**computation_result)
 

graphdatascience/procedure_surface/arrow/louvain_arrow_endpoints.py

@@ -50,12 +50,12 @@ def mutate(
             tolerance=tolerance,
         )
 
-        computation_result, node_properties_written = self._node_property_endpoints.run_job_and_mutate(
+        computation_result, mutate_result = self._node_property_endpoints.run_job_and_mutate(
             "v2/community.louvain", G, config, mutate_property
         )
 
-        computation_result["nodePropertiesWritten"] = node_properties_written
-        computation_result["mutateMillis"] = 0
+        computation_result["nodePropertiesWritten"] = mutate_result.node_properties_written
+        computation_result["mutateMillis"] = mutate_result.mutate_millis
 
         return LouvainMutateResult(**computation_result)
 

graphdatascience/procedure_surface/arrow/node_property_endpoints.py

@@ -3,6 +3,8 @@
 
 from pandas import DataFrame
 
+from graphdatascience.arrow_client.v2.api_types import MutateResult
+
 from ...arrow_client.authenticated_flight_client import AuthenticatedArrowClient
 from ...arrow_client.v2.data_mapper_utils import deserialize_single
 from ...arrow_client.v2.job_client import JobClient
@@ -30,12 +32,12 @@ def run_job_and_get_summary(self, endpoint: str, G: Graph, config: Dict[str, Any
 
     def run_job_and_mutate(
         self, endpoint: str, G: Graph, config: Dict[str, Any], mutate_property: str
-    ) -> tuple[Dict[str, Any], int]:
+    ) -> tuple[Dict[str, Any], MutateResult]:
         """Run a job, mutate node properties, and return summary with mutation result."""
         job_id = JobClient.run_job_and_wait(self._arrow_client, endpoint, config)
         mutate_result = MutationClient.mutate_node_property(self._arrow_client, job_id, mutate_property)
         computation_result = JobClient.get_summary(self._arrow_client, job_id)
-        return computation_result, mutate_result.node_properties_written
+        return computation_result, mutate_result
 
     def run_job_and_stream(self, endpoint: str, G: Graph, config: Dict[str, Any]) -> DataFrame:
         """Run a job and return streamed results."""

graphdatascience/procedure_surface/arrow/scc_arrow_endpoints.py

@@ -38,12 +38,12 @@ def mutate(
             sudo=sudo,
         )
 
-        computation_result, node_properties_written = self._node_property_endpoints.run_job_and_mutate(
+        computation_result, mutate_result = self._node_property_endpoints.run_job_and_mutate(
             "v2/community.scc", G, config, mutate_property
         )
 
-        computation_result["nodePropertiesWritten"] = node_properties_written
-        computation_result["mutateMillis"] = 0
+        computation_result["nodePropertiesWritten"] = mutate_result.node_properties_written
+        computation_result["mutateMillis"] = mutate_result.mutate_millis
 
         return SccMutateResult(**computation_result)
 

graphdatascience/procedure_surface/arrow/wcc_arrow_endpoints.py

@@ -44,12 +44,12 @@ def mutate(
             threshold=threshold,
         )
 
-        computation_result, node_properties_written = self._node_property_endpoints.run_job_and_mutate(
+        computation_result, mutate_result = self._node_property_endpoints.run_job_and_mutate(
             "v2/community.wcc", G, config, mutate_property
         )
 
-        computation_result["nodePropertiesWritten"] = node_properties_written
-        computation_result["mutateMillis"] = 0
+        computation_result["nodePropertiesWritten"] = mutate_result.node_properties_written
+        computation_result["mutateMillis"] = mutate_result.mutate_millis
 
         return WccMutateResult(**computation_result)