Share logic + move result classes

Author: FlorentinD

graphdatascience/arrow_client/v2/mutation_client.py

@@ -11,24 +11,41 @@ class MutationClient:
 
     @staticmethod
     def mutate_node_property(client: AuthenticatedArrowClient, job_id: str, mutate_property: str) -> MutateResult:
-        mutate_config = {"jobId": job_id, "mutateProperty": mutate_property}
-        start_time = time.time()
-        mutate_arrow_res = client.do_action_with_retry(MutationClient.MUTATE_ENDPOINT, mutate_config)
-        mutate_millis = math.ceil((time.time() - start_time) * 1000)
-        return MutateResult(mutateMillis=mutate_millis, **deserialize_single(mutate_arrow_res))
+        return MutationClient._mutate(
+            client=client,
+            job_id=job_id,
+            mutate_property=mutate_property,
+        )
 
     @staticmethod
     def mutate_relationship_property(
         client: AuthenticatedArrowClient,
         job_id: str,
         mutate_relationship_type: str,
         mutate_property: str,
+    ) -> MutateResult:
+        return MutationClient._mutate(
+            client=client,
+            job_id=job_id,
+            mutate_property=mutate_property,
+            mutate_relationship_type=mutate_relationship_type,
+        )
+
+    @staticmethod
+    def _mutate(
+        client: AuthenticatedArrowClient,
+        job_id: str,
+        mutate_property: str | None = None,
+        mutate_relationship_type: str | None = None,
     ) -> MutateResult:
         mutate_config = {
             "jobId": job_id,
-            "mutateProperty": mutate_property,
-            "mutateRelationshipType": mutate_relationship_type,
         }
+        if mutate_relationship_type:
+            mutate_config["mutateRelationshipType"] = mutate_relationship_type
+        if mutate_property:
+            mutate_config["mutateProperty"] = mutate_property
+
         start_time = time.time()
         mutate_arrow_res = client.do_action_with_retry(MutationClient.MUTATE_ENDPOINT, mutate_config)
         mutate_millis = math.ceil((time.time() - start_time) * 1000)

graphdatascience/procedure_surface/api/similarity/knn_endpoints.py

@@ -10,53 +10,6 @@
 from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
 
 
-class KnnMutateResult(BaseResult):
-    """Represents the result of running K-Nearest Neighbors in mutate mode."""
-
-    pre_processing_millis: int
-    compute_millis: int
-    mutate_millis: int
-    post_processing_millis: int
-    nodes_compared: int
-    relationships_written: int
-    similarity_distribution: dict[str, Any]
-    did_converge: bool
-    ran_iterations: int
-    node_pairs_considered: int
-    configuration: dict[str, Any]
-
-
-class KnnStatsResult(BaseResult):
-    """Represents the result of running K-Nearest Neighbors in stats mode."""
-
-    pre_processing_millis: int
-    compute_millis: int
-    post_processing_millis: int
-    nodes_compared: int
-    similarity_pairs: int
-    similarity_distribution: dict[str, Any]
-    did_converge: bool
-    ran_iterations: int
-    node_pairs_considered: int
-    configuration: dict[str, Any]
-
-
-class KnnWriteResult(BaseResult):
-    """Represents the result of running K-Nearest Neighbors in write mode."""
-
-    pre_processing_millis: int
-    compute_millis: int
-    write_millis: int
-    post_processing_millis: int
-    nodes_compared: int
-    relationships_written: int
-    did_converge: bool
-    ran_iterations: int
-    node_pairs_considered: int
-    similarity_distribution: dict[str, Any]
-    configuration: dict[str, Any]
-
-
 class KnnEndpoints(ABC):
     @abstractmethod
     def mutate(
@@ -85,8 +38,6 @@ def mutate(
         """
         Runs the K-Nearest Neighbors algorithm and stores the results as new relationships in the graph catalog.
 
-        The K-Nearest Neighbors algorithm computes a distance value for all node pairs in the graph and creates new relationships between each node and its k nearest neighbors
-
         Parameters
         ----------
         G : GraphV2
@@ -161,8 +112,6 @@ def stats(
         """
         Runs the K-Nearest Neighbors algorithm and returns execution statistics.
 
-        The K-Nearest Neighbors algorithm computes a distance value for all node pairs in the graph and creates new relationships between each node and its k nearest neighbors
-
         Parameters
         ----------
         G : GraphV2
@@ -233,8 +182,6 @@ def stream(
         """
         Runs the K-Nearest Neighbors algorithm and returns the result as a DataFrame.
 
-        The K-Nearest Neighbors algorithm computes a distance value for all node pairs in the graph and creates new relationships between each node and its k nearest neighbors
-
         Parameters
         ----------
         G : GraphV2
@@ -308,8 +255,6 @@ def write(
         """
         Runs the K-Nearest Neighbors algorithm and writes the results back to the database.
 
-        The K-Nearest Neighbors algorithm computes a distance value for all node pairs in the graph and creates new relationships between each node and its k nearest neighbors
-
         Parameters
         ----------
         G : GraphV2
@@ -386,8 +331,6 @@ def estimate(
         """
         Estimates the memory requirements for running the K-Nearest Neighbors algorithm.
 
-        The K-Nearest Neighbors algorithm computes a distance value for all node pairs in the graph and creates new relationships between each node and its k nearest neighbors
-
         Parameters
         ----------
         G : GraphV2
@@ -432,3 +375,44 @@ def estimate(
         EstimationResult
             Object containing the estimated memory requirements.
         """
+
+
+class KnnMutateResult(BaseResult):
+    pre_processing_millis: int
+    compute_millis: int
+    mutate_millis: int
+    post_processing_millis: int
+    nodes_compared: int
+    relationships_written: int
+    similarity_distribution: dict[str, Any]
+    did_converge: bool
+    ran_iterations: int
+    node_pairs_considered: int
+    configuration: dict[str, Any]
+
+
+class KnnStatsResult(BaseResult):
+    pre_processing_millis: int
+    compute_millis: int
+    post_processing_millis: int
+    nodes_compared: int
+    similarity_pairs: int
+    similarity_distribution: dict[str, Any]
+    did_converge: bool
+    ran_iterations: int
+    node_pairs_considered: int
+    configuration: dict[str, Any]
+
+
+class KnnWriteResult(BaseResult):
+    pre_processing_millis: int
+    compute_millis: int
+    write_millis: int
+    post_processing_millis: int
+    nodes_compared: int
+    relationships_written: int
+    did_converge: bool
+    ran_iterations: int
+    node_pairs_considered: int
+    similarity_distribution: dict[str, Any]
+    configuration: dict[str, Any]