Add implementation for Wcc.estimate for Arrow endpoints

Author: DarthMax

graphdatascience/procedure_surface/api/wcc_endpoints.py

@@ -246,7 +246,7 @@ def write(
     @abstractmethod
     def estimate(
         self,
-        graph_name: Optional[str] = None,
+        G: Optional[Graph] = None,
         projection_config: Optional[dict[str, Any]] = None,
     ) -> EstimationResult:
         """
@@ -259,8 +259,8 @@ def estimate(
 
         Parameters
         ----------
-        graph_name : Optional[str], optional
-            Name of the graph to be used in the estimation
+        G : Optional[Graph], optional
+            The graph to be used in the estimation
         projection_config : Optional[dict[str, Any]], optional
             Configuration dictionary for the projection.
 

graphdatascience/procedure_surface/arrow/arrow_wcc_endpoints.py renamed to graphdatascience/procedure_surface/arrow/wcc_arrow_endpoints.py

@@ -1,8 +1,10 @@
+import json
 from typing import Any, List, Optional
 
 from pandas import DataFrame
 
 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
 from ...arrow_client.v2.mutation_client import MutationClient
 from ...arrow_client.v2.write_back_client import WriteBackClient
@@ -175,6 +177,15 @@ def write(
         return WccWriteResult(**computation_result)
 
     def estimate(
-        self, graph_name: Optional[str] = None, projection_config: Optional[dict[str, Any]] = None
+        self, G: Optional[Graph] = None, projection_config: Optional[dict[str, Any]] = None
     ) -> EstimationResult:
-        pass
+        if G is not None:
+            payload = {"graphName": G.name()}
+        elif projection_config is not None:
+            payload = projection_config
+        else:
+            raise ValueError("Either graph_name or projection_config must be provided.")
+
+        res = self._arrow_client.do_action_with_retry("v2/community.wcc.estimate", json.dumps(payload).encode("utf-8"))
+
+        return EstimationResult(**deserialize_single(res))

graphdatascience/procedure_surface/cypher/wcc_cypher_endpoints.py

@@ -1,3 +1,4 @@
+from collections import OrderedDict
 from typing import Any, List, Optional, Union
 
 from pandas import DataFrame
@@ -177,18 +178,20 @@ def write(
         return WccWriteResult(**result.to_dict())
 
     def estimate(
-        self, graph_name: Optional[str] = None, projection_config: Optional[dict[str, Any]] = None
+        self, G: Optional[Graph] = None, projection_config: Optional[dict[str, Any]] = None
     ) -> EstimationResult:
-        config: Union[str, dict[str, Any]] = {}
+        config: Union[dict[str, Any]] = OrderedDict()
 
-        if graph_name is not None:
-            config = graph_name
+        if G is not None:
+            config["graphNameOrConfiguration"] = G.name()
         elif projection_config is not None:
-            config = projection_config
+            config["graphNameOrConfiguration"] = projection_config
         else:
             raise ValueError("Either graph_name or projection_config must be provided.")
 
-        params = CallParameters(config=config)
+        config["algoConfig"] = {}
+
+        params = CallParameters(**config)
 
         result = self._query_runner.call_procedure(endpoint="gds.wcc.stats.estimate", params=params).squeeze()
 

graphdatascience/tests/integrationV2/procedure_surface/arrow/conftest.py

@@ -5,7 +5,6 @@
 import pytest
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.waiting_utils import wait_for_logs
-from testcontainers.neo4j import Neo4jContainer
 
 from graphdatascience.arrow_client.arrow_authentication import UsernamePasswordAuthentication
 from graphdatascience.arrow_client.arrow_info import ArrowInfo
@@ -28,27 +27,12 @@ def password_file() -> Generator[str, None, None]:
     os.rmdir(temp_dir)
 
 
-@pytest.fixture(scope="session")
-def neo4j_database_container() -> Generator[Neo4jContainer, None, None]:
-    neo4j_image = os.getenv("NEO4J_DATABASE_IMAGE", "neo4j:5.11-enterprise")
-
-    neo4j_container = (
-        Neo4jContainer(
-            image=neo4j_image,
-        )
-        .with_env("NEO4J_ACCEPT_LICENSE_AGREEMENT", "yes")
-        .with_env("NEO4J_dbms_security_procedures_unrestricted", "gds.*")
-        .with_env("NEO4J_dbms_security_procedures_allowlist", "gds.*")
-    )
-
-    with neo4j_container as neo4j_db:
-        wait_for_logs(neo4j_db)
-        yield neo4j_db
-
-
 @pytest.fixture(scope="session")
 def session_container(password_file: str) -> Generator[DockerContainer, None, None]:
-    session_image = os.getenv("GDS_SESSION_IMAGE")
+    session_image = os.getenv(
+        "GDS_SESSION_IMAGE",
+        "europe-west1-docker.pkg.dev/aura-docker-images/aura/gds-session:97ac47f7928c0533b3099539f0f5b3058a52c203",
+    )
 
     if session_image is None:
         raise ValueError("GDS_SESSION_IMAGE environment variable is not set")

graphdatascience/tests/integrationV2/procedure_surface/arrow/test_wcc_arrow_endpoints.py

@@ -5,7 +5,7 @@
 
 from graphdatascience import Graph
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
-from graphdatascience.procedure_surface.arrow.arrow_wcc_endpoints import WccArrowEndpoints
+from graphdatascience.procedure_surface.arrow.wcc_arrow_endpoints import WccArrowEndpoints
 
 
 class MockGraph(Graph):
@@ -71,3 +71,17 @@ def test_wcc_mutate(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph) -> No
     assert result.post_processing_millis >= 0
     assert result.mutate_millis >= 0
     assert result.node_properties_written == 3
+
+
+def test_wcc_estimate(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph) -> None:
+    result = wcc_endpoints.estimate(sample_graph)
+
+    assert result.node_count == 3
+    assert result.relationship_count == 1
+    assert "Bytes" in result.required_memory
+    # assert result.tree_view > 0
+    # assert result.map_view > 0
+    assert result.bytes_min > 0
+    assert result.bytes_max > 0
+    assert result.heap_percentage_min > 0
+    assert result.heap_percentage_max > 0

graphdatascience/tests/integrationV2/procedure_surface/cypher/conftest.py

@@ -11,7 +11,7 @@
 
 
 @pytest.fixture(scope="session")
-def neo4j_database_container() -> Generator[Neo4jContainer, None, None]:
+def gds_plugin_container() -> Generator[Neo4jContainer, None, None]:
     neo4j_image = os.getenv("NEO4J_DATABASE_IMAGE", "neo4j:enterprise")
 
     neo4j_container = (
@@ -28,8 +28,8 @@ def neo4j_database_container() -> Generator[Neo4jContainer, None, None]:
 
 
 @pytest.fixture
-def query_runner(neo4j_database_container: DockerContainer) -> Generator[QueryRunner, None, None]:
+def query_runner(gds_plugin_container: DockerContainer) -> Generator[QueryRunner, None, None]:
     yield Neo4jQueryRunner.create_for_db(
-        f"bolt://localhost:{neo4j_database_container.get_exposed_port(7687)}",
+        f"bolt://localhost:{gds_plugin_container.get_exposed_port(7687)}",
         ("neo4j", "password"),
     )

graphdatascience/tests/integrationV2/procedure_surface/cypher/test_wcc_cypher_endpoints.py

@@ -3,7 +3,7 @@
 import pytest
 
 from graphdatascience import Graph, QueryRunner
-from graphdatascience.procedure_surface.arrow.arrow_wcc_endpoints import WccArrowEndpoints
+from graphdatascience.procedure_surface.arrow.wcc_arrow_endpoints import WccArrowEndpoints
 from graphdatascience.procedure_surface.cypher.wcc_cypher_endpoints import WccCypherEndpoints
 
 
@@ -73,3 +73,17 @@ def test_wcc_mutate(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph) -> No
     assert result.post_processing_millis >= 0
     assert result.mutate_millis >= 0
     assert result.node_properties_written == 3
+
+
+def test_wcc_estimate(wcc_endpoints: WccArrowEndpoints, sample_graph: Graph) -> None:
+    result = wcc_endpoints.estimate(sample_graph)
+
+    assert result.node_count == 3
+    assert result.relationship_count == 1
+    assert "Bytes" in result.required_memory
+    # assert result.tree_view > 0
+    # assert result.map_view > 0
+    assert result.bytes_min > 0
+    assert result.bytes_max > 0
+    assert result.heap_percentage_min > 0
+    assert result.heap_percentage_max > 0

graphdatascience/tests/unit/procedure_surface/cypher/test_wcc_cypher_endpoints.py

@@ -337,7 +337,7 @@ def test_estimate_with_graph_name(graph: Graph) -> None:
 
     query_runner = CollectingQueryRunner(DEFAULT_SERVER_VERSION, {"wcc.stats.estimate": pd.DataFrame([result])})
 
-    estimate = WccCypherEndpoints(query_runner).estimate(graph_name=graph.name())
+    estimate = WccCypherEndpoints(query_runner).estimate(G=graph)
 
     assert estimate.node_count == 100
     assert estimate.relationship_count == 200
@@ -348,7 +348,7 @@ def test_estimate_with_graph_name(graph: Graph) -> None:
     assert len(query_runner.queries) == 1
     assert "gds.wcc.stats.estimate" in query_runner.queries[0]
     params = query_runner.params[0]
-    assert params["config"] == "test_graph"
+    assert params["graphNameOrConfiguration"] == "test_graph"
 
 
 def test_estimate_with_projection_config() -> None: