Implement WccEndpoints with cypher

Author: DarthMax

graphdatascience/procedure_surface/cypher/wcc_cypher_endpoints.py

@@ -0,0 +1,195 @@
+from typing import Any, List, Optional, Union
+
+from pandas import DataFrame
+
+from ...call_parameters import CallParameters
+from ...graph.graph_object import Graph
+from ...query_runner.query_runner import QueryRunner
+from ..api.estimation_result import EstimationResult
+from ..api.wcc_endpoints import WccEndpoints, WccMutateResult, WccStatsResult, WccWriteResult
+from ..utils.config_converter import ConfigConverter
+
+
+class WccCypherEndpoints(WccEndpoints):
+    """
+    Implementation of the WCC algorithm endpoints.
+    This class handles the actual execution by forwarding calls to the query runner.
+    """
+
+    def __init__(self, query_runner: QueryRunner):
+        self._query_runner = query_runner
+
+    def mutate(
+        self,
+        G: Graph,
+        mutate_property: str,
+        threshold: Optional[float] = None,
+        relationship_types: Optional[List[str]] = None,
+        node_labels: Optional[List[str]] = None,
+        sudo: Optional[bool] = None,
+        log_progress: Optional[bool] = None,
+        username: Optional[str] = None,
+        concurrency: Optional[int] = None,
+        job_id: Optional[str] = None,
+        seed_property: Optional[str] = None,
+        consecutive_ids: Optional[bool] = None,
+        relationship_weight_property: Optional[str] = None,
+    ) -> WccMutateResult:
+        config = ConfigConverter.convert_to_gds_config(
+            mutate_property=mutate_property,
+            concurrency=concurrency,
+            consecutive_ids=consecutive_ids,
+            job_id=job_id,
+            log_progress=log_progress,
+            node_labels=node_labels,
+            relationship_types=relationship_types,
+            relationship_weight_property=relationship_weight_property,
+            seed_property=seed_property,
+            sudo=sudo,
+            threshold=threshold,
+            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.mutate", params=params).squeeze()
+
+        return WccMutateResult(**cypher_result.to_dict())
+
+    def stats(
+        self,
+        G: Graph,
+        threshold: Optional[float] = None,
+        relationship_types: Optional[List[str]] = None,
+        node_labels: Optional[List[str]] = None,
+        sudo: Optional[bool] = None,
+        log_progress: Optional[bool] = None,
+        username: Optional[str] = None,
+        concurrency: Optional[int] = None,
+        job_id: Optional[str] = None,
+        seed_property: Optional[str] = None,
+        consecutive_ids: Optional[bool] = None,
+        relationship_weight_property: Optional[str] = None,
+    ) -> WccStatsResult:
+        config = ConfigConverter.convert_to_gds_config(
+            concurrency=concurrency,
+            consecutive_ids=consecutive_ids,
+            job_id=job_id,
+            log_progress=log_progress,
+            node_labels=node_labels,
+            relationship_types=relationship_types,
+            relationship_weight_property=relationship_weight_property,
+            seed_property=seed_property,
+            sudo=sudo,
+            threshold=threshold,
+            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
+
+        return WccStatsResult(**cypher_result.to_dict())
+
+    def stream(
+        self,
+        G: Graph,
+        min_component_size: Optional[int] = None,
+        threshold: Optional[float] = None,
+        relationship_types: Optional[List[str]] = None,
+        node_labels: Optional[List[str]] = None,
+        sudo: Optional[bool] = None,
+        log_progress: Optional[bool] = None,
+        username: Optional[str] = None,
+        concurrency: Optional[int] = None,
+        job_id: Optional[str] = None,
+        seed_property: Optional[str] = None,
+        consecutive_ids: Optional[bool] = None,
+        relationship_weight_property: Optional[str] = None,
+    ) -> DataFrame:
+        config = ConfigConverter.convert_to_gds_config(
+            concurrency=concurrency,
+            consecutive_ids=consecutive_ids,
+            job_id=job_id,
+            log_progress=log_progress,
+            min_component_size=min_component_size,
+            node_labels=node_labels,
+            relationship_types=relationship_types,
+            relationship_weight_property=relationship_weight_property,
+            seed_property=seed_property,
+            sudo=sudo,
+            threshold=threshold,
+            username=username,
+        )
+
+        # Run procedure and return results
+        params = CallParameters(graph_name=G.name(), config=config)
+        params.ensure_job_id_in_config()
+
+        return self._query_runner.call_procedure(endpoint="gds.wcc.stream", params=params)
+
+    def write(
+        self,
+        G: Graph,
+        write_property: str,
+        min_component_size: Optional[int] = None,
+        threshold: Optional[float] = None,
+        relationship_types: Optional[List[str]] = None,
+        node_labels: Optional[List[str]] = None,
+        sudo: Optional[bool] = None,
+        log_progress: Optional[bool] = None,
+        username: Optional[str] = None,
+        concurrency: Optional[int] = None,
+        job_id: Optional[str] = None,
+        seed_property: Optional[str] = None,
+        consecutive_ids: Optional[bool] = None,
+        relationship_weight_property: Optional[str] = None,
+        write_concurrency: Optional[int] = None,
+    ) -> WccWriteResult:
+        config = ConfigConverter.convert_to_gds_config(
+            write_property=write_property,
+            concurrency=concurrency,
+            consecutive_ids=consecutive_ids,
+            job_id=job_id,
+            log_progress=log_progress,
+            min_component_size=min_component_size,
+            node_labels=node_labels,
+            relationship_types=relationship_types,
+            relationship_weight_property=relationship_weight_property,
+            seed_property=seed_property,
+            sudo=sudo,
+            threshold=threshold,
+            username=username,
+        )
+
+        if write_concurrency is not None:
+            config["writeConcurrency"] = write_concurrency
+
+        params = CallParameters(graph_name=G.name(), config=config)
+        params.ensure_job_id_in_config()
+
+        result = self._query_runner.call_procedure(endpoint="gds.wcc.write", params=params).squeeze()  # type: ignore
+
+        return WccWriteResult(**result.to_dict())
+
+    def estimate(
+        self, graph_name: Optional[str] = None, projection_config: Optional[dict[str, Any]] = None
+    ) -> EstimationResult:
+        config: Union[str, dict[str, Any]] = {}
+
+        if graph_name is not None:
+            config = graph_name
+        elif projection_config is not None:
+            config = projection_config
+        else:
+            raise ValueError("Either graph_name or projection_config must be provided.")
+
+        params = CallParameters(config=config)
+
+        result = self._query_runner.call_procedure(endpoint="gds.wcc.stats.estimate", params=params).squeeze()
+
+        return EstimationResult(**result.to_dict())

graphdatascience/procedure_surface/utils/config_converter.py

@@ -0,0 +1,33 @@
+from typing import Any, Dict, Optional
+
+
+class ConfigConverter:
+    @staticmethod
+    def convert_to_gds_config(**kwargs: Optional[Any]) -> dict[str, Any]:
+        config: dict[str, Any] = {}
+
+        # Process kwargs
+        processed_kwargs = ConfigConverter._process_dict_values(kwargs)
+        config.update(processed_kwargs)
+
+        return config
+
+    @staticmethod
+    def _convert_to_camel_case(name: str) -> str:
+        """Convert a snake_case string to camelCase."""
+        parts = name.split("_")
+        return "".join([word.capitalize() if i > 0 else word.lower() for i, word in enumerate(parts)])
+
+    @staticmethod
+    def _process_dict_values(input_dict: Dict[str, Any]) -> Dict[str, Any]:
+        """Process dictionary values, converting keys to camelCase and handling nested dictionaries."""
+        result = {}
+        for key, value in input_dict.items():
+            if value is not None:
+                camel_key = ConfigConverter._convert_to_camel_case(key)
+                # Recursively process nested dictionaries
+                if isinstance(value, dict):
+                    result[camel_key] = ConfigConverter._process_dict_values(value)
+                else:
+                    result[camel_key] = value
+        return result

graphdatascience/tests/unit/procedure_surface/__init__.py

@@ -0,0 +1,72 @@
+from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
+
+
+def test_estimation_result_initialization() -> None:
+    estimation_result = EstimationResult(
+        nodeCount=5,
+        relationshipCount=10,
+        requiredMemory="512MB",
+        treeView="TreeData",
+        mapView={"key": "value"},
+        bytesMin=500,
+        bytesMax=1000,
+        heapPercentageMin=0.1,
+        heapPercentageMax=0.5,
+    )
+    assert estimation_result.node_count == 5
+    assert estimation_result.relationship_count == 10
+    assert estimation_result.required_memory == "512MB"
+    assert estimation_result.tree_view == "TreeData"
+    assert estimation_result.map_view == {"key": "value"}
+    assert estimation_result.bytes_min == 500
+    assert estimation_result.bytes_max == 1000
+    assert estimation_result.heap_percentage_min == 0.1
+    assert estimation_result.heap_percentage_max == 0.5
+
+
+def test_estimation_result_getitem() -> None:
+    estimation_result = EstimationResult(
+        nodeCount=5,
+        relationshipCount=10,
+        requiredMemory="512MB",
+        treeView="TreeData",
+        mapView={"key": "value"},
+        bytesMin=500,
+        bytesMax=1000,
+        heapPercentageMin=0.1,
+        heapPercentageMax=0.5,
+    )
+    assert estimation_result["node_count"] == 5
+    assert estimation_result["relationship_count"] == 10
+    assert estimation_result["required_memory"] == "512MB"
+    assert estimation_result["tree_view"] == "TreeData"
+    assert estimation_result["map_view"] == {"key": "value"}
+    assert estimation_result["bytes_min"] == 500
+    assert estimation_result["bytes_max"] == 1000
+    assert estimation_result["heap_percentage_min"] == 0.1
+    assert estimation_result["heap_percentage_max"] == 0.5
+
+
+def test_estimation_result_from_cypher() -> None:
+    cypher_result = {
+        "nodeCount": 5,
+        "relationshipCount": 10,
+        "requiredMemory": "512MB",
+        "treeView": "TreeData",
+        "mapView": {"key": "value"},
+        "bytesMin": 500,
+        "bytesMax": 1000,
+        "heapPercentageMin": 0.1,
+        "heapPercentageMax": 0.5,
+    }
+    estimation_result = EstimationResult.from_cypher(cypher_result)
+    assert isinstance(estimation_result, EstimationResult)
+    assert estimation_result.node_count == 5
+    assert estimation_result.relationship_count == 10
+    assert estimation_result.required_memory == "512MB"
+    assert estimation_result.tree_view == "TreeData"
+    assert estimation_result.map_view == {"key": "value"}
+    assert estimation_result.bytes_min == 500
+    assert estimation_result.bytes_max == 1000
+    assert estimation_result.heap_percentage_min == 0.1
+    assert estimation_result.heap_percentage_max == 0.5