Extract common method cypher estimation

Co-authored-by: Florentin Dörre <florentin.dorre@neo4j.com>

Author: DarthMax

graphdatascience/procedure_surface/cypher/__init__.py

@@ -0,0 +1,58 @@
+from collections import OrderedDict
+from typing import Any, Optional, Union
+
+from graphdatascience.call_parameters import CallParameters
+from graphdatascience.graph.graph_object import Graph
+from graphdatascience.query_runner.query_runner import QueryRunner
+from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
+
+
+def estimate_algorithm(
+    endpoint: str,
+    query_runner: QueryRunner,
+    G: Optional[Graph] = None,
+    projection_config: Optional[dict[str, Any]] = None,
+) -> EstimationResult:
+    """
+    Estimate the memory consumption of an algorithm run.
+
+    This utility function provides a common implementation for estimation
+    across all cypher endpoint implementations.
+
+    Parameters
+    ----------
+    query_runner : QueryRunner
+        The query runner to use for the estimation call
+    endpoint : str
+        The full endpoint name for the estimation procedure (e.g., "gds.kcore.stats.estimate")
+    G : Optional[Graph], optional
+        The graph to be used in the estimation
+    projection_config : Optional[dict[str, Any]], optional
+        Configuration dictionary for the projection
+
+    Returns
+    -------
+    EstimationResult
+        An object containing the result of the estimation
+
+    Raises
+    ------
+    ValueError
+        If neither G nor projection_config is provided
+    """
+    config: Union[dict[str, Any]] = OrderedDict()
+
+    if G is not None:
+        config["graphNameOrConfiguration"] = G.name()
+    elif projection_config is not None:
+        config["graphNameOrConfiguration"] = projection_config
+    else:
+        raise ValueError("Either graph_name or projection_config must be provided.")
+
+    config["algoConfig"] = {}
+
+    params = CallParameters(**config)
+
+    result = query_runner.call_procedure(endpoint=endpoint, params=params).squeeze()
+
+    return EstimationResult(**result.to_dict())

graphdatascience/procedure_surface/cypher/estimation_utils.py

@@ -1,5 +1,4 @@
-from collections import OrderedDict
-from typing import Any, List, Optional, Union
+from typing import Any, List, Optional
 
 from pandas import DataFrame
 
@@ -14,6 +13,7 @@
     K1ColoringWriteResult,
 )
 from ..utils.config_converter import ConfigConverter
+from graphdatascience.procedure_surface.cypher.estimation_utils import estimate_algorithm
 
 
 class K1ColoringCypherEndpoints(K1ColoringEndpoints):
@@ -169,19 +169,9 @@ def write(
     def estimate(
         self, G: Optional[Graph] = None, projection_config: Optional[dict[str, Any]] = None
     ) -> EstimationResult:
-        config: Union[dict[str, Any]] = OrderedDict()
-
-        if G is not None:
-            config["graphNameOrConfiguration"] = G.name()
-        elif projection_config is not None:
-            config["graphNameOrConfiguration"] = projection_config
-        else:
-            raise ValueError("Either graph_name or projection_config must be provided.")
-
-        config["algoConfig"] = {}
-
-        params = CallParameters(**config)
-
-        result = self._query_runner.call_procedure(endpoint="gds.k1coloring.stats.estimate", params=params).squeeze()
-
-        return EstimationResult(**result.to_dict())
+        return estimate_algorithm(
+            endpoint="gds.k1coloring.stats.estimate",
+            query_runner=self._query_runner,
+            G=G,
+            projection_config=projection_config,
+        )

graphdatascience/procedure_surface/cypher/k1coloring_cypher_endpoints.py

@@ -1,5 +1,4 @@
-from collections import OrderedDict
-from typing import Any, List, Optional, Union
+from typing import Any, List, Optional
 
 from pandas import DataFrame
 
@@ -9,6 +8,7 @@
 from ..api.estimation_result import EstimationResult
 from ..api.kcore_endpoints import KCoreEndpoints, KCoreMutateResult, KCoreStatsResult, KCoreWriteResult
 from ..utils.config_converter import ConfigConverter
+from graphdatascience.procedure_surface.cypher.estimation_utils import estimate_algorithm
 
 
 class KCoreCypherEndpoints(KCoreEndpoints):
@@ -145,19 +145,9 @@ def write(
     def estimate(
         self, G: Optional[Graph] = None, projection_config: Optional[dict[str, Any]] = None
     ) -> EstimationResult:
-        config: Union[dict[str, Any]] = OrderedDict()
-
-        if G is not None:
-            config["graphNameOrConfiguration"] = G.name()
-        elif projection_config is not None:
-            config["graphNameOrConfiguration"] = projection_config
-        else:
-            raise ValueError("Either graph_name or projection_config must be provided.")
-
-        config["algoConfig"] = {}
-
-        params = CallParameters(**config)
-
-        result = self._query_runner.call_procedure(endpoint="gds.kcore.stats.estimate", params=params).squeeze()
-
-        return EstimationResult(**result.to_dict())
+        return estimate_algorithm(
+            endpoint="gds.kcore.stats.estimate",
+            query_runner=self._query_runner,
+            G=G,
+            projection_config=projection_config,
+        )

graphdatascience/procedure_surface/cypher/kcore_cypher_endpoints.py

@@ -1,5 +1,4 @@
-from collections import OrderedDict
-from typing import Any, List, Optional, Union
+from typing import Any, List, Optional
 
 from pandas import DataFrame
 
@@ -9,6 +8,7 @@
 from ..api.estimation_result import EstimationResult
 from ..api.louvain_endpoints import LouvainEndpoints, LouvainMutateResult, LouvainStatsResult, LouvainWriteResult
 from ..utils.config_converter import ConfigConverter
+from graphdatascience.procedure_surface.cypher.estimation_utils import estimate_algorithm
 
 
 class LouvainCypherEndpoints(LouvainEndpoints):
@@ -57,7 +57,6 @@ def mutate(
             username=username,
         )
 
-        # Run procedure and return results
         params = CallParameters(graph_name=G.name(), config=config)
         params.ensure_job_id_in_config()
 
@@ -100,11 +99,10 @@ def stats(
             username=username,
         )
 
-        # Run procedure and return results
         params = CallParameters(graph_name=G.name(), config=config)
         params.ensure_job_id_in_config()
 
-        cypher_result = self._query_runner.call_procedure(endpoint="gds.louvain.stats", params=params).squeeze()  # type: ignore
+        cypher_result = self._query_runner.call_procedure(endpoint="gds.louvain.stats", params=params).squeeze()
 
         return LouvainStatsResult(**cypher_result.to_dict())
 
@@ -145,7 +143,6 @@ def stream(
             username=username,
         )
 
-        # Run procedure and return results
         params = CallParameters(graph_name=G.name(), config=config)
         params.ensure_job_id_in_config()
 
@@ -206,19 +203,9 @@ def write(
     def estimate(
         self, G: Optional[Graph] = None, projection_config: Optional[dict[str, Any]] = None
     ) -> EstimationResult:
-        config: Union[dict[str, Any]] = OrderedDict()
-
-        if G is not None:
-            config["graphNameOrConfiguration"] = G.name()
-        elif projection_config is not None:
-            config["graphNameOrConfiguration"] = projection_config
-        else:
-            raise ValueError("Either graph_name or projection_config must be provided.")
-
-        config["algoConfig"] = {}
-
-        params = CallParameters(**config)
-
-        result = self._query_runner.call_procedure(endpoint="gds.louvain.stats.estimate", params=params).squeeze()
-
-        return EstimationResult(**result.to_dict())
+        return estimate_algorithm(
+            endpoint="gds.louvain.stats.estimate",
+            query_runner=self._query_runner,
+            G=G,
+            projection_config=projection_config,
+        )

graphdatascience/procedure_surface/cypher/louvain_cypher_endpoints.py

@@ -1,5 +1,4 @@
-from collections import OrderedDict
-from typing import Any, List, Optional, Union
+from typing import Any, List, Optional
 
 from pandas import DataFrame
 
@@ -9,11 +8,12 @@
 from ..api.estimation_result import EstimationResult
 from ..api.scc_endpoints import SccEndpoints, SccMutateResult, SccStatsResult, SccWriteResult
 from ..utils.config_converter import ConfigConverter
+from graphdatascience.procedure_surface.cypher.estimation_utils import estimate_algorithm
 
 
 class SccCypherEndpoints(SccEndpoints):
     """
-    Implementation of the SCC algorithm endpoints.
+    Implementation of the Strongly Connected Components (SCC) algorithm endpoints.
     This class handles the actual execution by forwarding calls to the query runner.
     """
 
@@ -153,19 +153,6 @@ def write(
     def estimate(
         self, G: Optional[Graph] = None, projection_config: Optional[dict[str, Any]] = None
     ) -> EstimationResult:
-        config: Union[dict[str, Any]] = OrderedDict()
-
-        if G is not None:
-            config["graphNameOrConfiguration"] = G.name()
-        elif projection_config is not None:
-            config["graphNameOrConfiguration"] = projection_config
-        else:
-            raise ValueError("Either graph_name or projection_config must be provided.")
-
-        config["algoConfig"] = {}
-
-        params = CallParameters(**config)
-
-        result = self._query_runner.call_procedure(endpoint="gds.scc.stats.estimate", params=params).squeeze()
-
-        return EstimationResult(**result.to_dict())
+        return estimate_algorithm(
+            endpoint="gds.scc.stats.estimate", query_runner=self._query_runner, G=G, projection_config=projection_config
+        )

graphdatascience/procedure_surface/cypher/scc_cypher_endpoints.py

@@ -1,5 +1,4 @@
-from collections import OrderedDict
-from typing import Any, List, Optional, Union
+from typing import Any, List, Optional
 
 from pandas import DataFrame
 
@@ -9,11 +8,12 @@
 from ..api.estimation_result import EstimationResult
 from ..api.wcc_endpoints import WccEndpoints, WccMutateResult, WccStatsResult, WccWriteResult
 from ..utils.config_converter import ConfigConverter
+from graphdatascience.procedure_surface.cypher.estimation_utils import estimate_algorithm
 
 
 class WccCypherEndpoints(WccEndpoints):
     """
-    Implementation of the WCC algorithm endpoints.
+    Implementation of the Weakly Connected Components (WCC) algorithm endpoints.
     This class handles the actual execution by forwarding calls to the query runner.
     """
 
@@ -51,7 +51,6 @@ def mutate(
             username=username,
         )
 
-        # Run procedure and return results
         params = CallParameters(graph_name=G.name(), config=config)
         params.ensure_job_id_in_config()
 
@@ -88,11 +87,10 @@ def stats(
             username=username,
         )
 
-        # Run procedure and return results
         params = CallParameters(graph_name=G.name(), config=config)
         params.ensure_job_id_in_config()
 
-        cypher_result = self._query_runner.call_procedure(endpoint="gds.wcc.stats", params=params).squeeze()  # type: ignore
+        cypher_result = self._query_runner.call_procedure(endpoint="gds.wcc.stats", params=params).squeeze()
 
         return WccStatsResult(**cypher_result.to_dict())
 
@@ -127,7 +125,6 @@ def stream(
             username=username,
         )
 
-        # Run procedure and return results
         params = CallParameters(graph_name=G.name(), config=config)
         params.ensure_job_id_in_config()
 
@@ -180,19 +177,6 @@ def write(
     def estimate(
         self, G: Optional[Graph] = None, projection_config: Optional[dict[str, Any]] = None
     ) -> EstimationResult:
-        config: Union[dict[str, Any]] = OrderedDict()
-
-        if G is not None:
-            config["graphNameOrConfiguration"] = G.name()
-        elif projection_config is not None:
-            config["graphNameOrConfiguration"] = projection_config
-        else:
-            raise ValueError("Either graph_name or projection_config must be provided.")
-
-        config["algoConfig"] = {}
-
-        params = CallParameters(**config)
-
-        result = self._query_runner.call_procedure(endpoint="gds.wcc.stats.estimate", params=params).squeeze()
-
-        return EstimationResult(**result.to_dict())
+        return estimate_algorithm(
+            endpoint="gds.wcc.stats.estimate", query_runner=self._query_runner, G=G, projection_config=projection_config
+        )