Move filtered knn behind knn endpoints

FlorentinD · FlorentinD · commit 186b4808f57e · 2025-10-17T10:17:47.000+02:00
diff --git a/graphdatascience/procedure_surface/api/similarity/knn_endpoints.py b/graphdatascience/procedure_surface/api/similarity/knn_endpoints.py
@@ -5,12 +5,21 @@
 
 from pandas import DataFrame
 
-from graphdatascience.procedure_surface.api.base_result import BaseResult
 from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
 from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
+from graphdatascience.procedure_surface.api.similarity.knn_filtered_endpoints import KnnFilteredEndpoints
+from graphdatascience.procedure_surface.api.similarity.knn_results import (
+    KnnMutateResult,
+    KnnStatsResult,
+    KnnWriteResult,
+)
 
 
 class KnnEndpoints(ABC):
+    @abstractmethod
+    def filtered(self) -> KnnFilteredEndpoints:
+        pass
+
     @abstractmethod
     def mutate(
         self,
@@ -374,44 +383,3 @@ 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, int | float]
-    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, int | float]
-    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, int | float]
-    configuration: dict[str, Any]
diff --git a/graphdatascience/procedure_surface/api/similarity/knn_filtered_endpoints.py b/graphdatascience/procedure_surface/api/similarity/knn_filtered_endpoints.py
@@ -7,7 +7,7 @@
 
 from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
 from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
-from graphdatascience.procedure_surface.api.similarity.knn_endpoints import (
+from graphdatascience.procedure_surface.api.similarity.knn_results import (
     KnnMutateResult,
     KnnStatsResult,
     KnnWriteResult,
diff --git a/graphdatascience/procedure_surface/api/similarity/knn_results.py b/graphdatascience/procedure_surface/api/similarity/knn_results.py
@@ -0,0 +1,44 @@
+from typing import Any
+
+from graphdatascience.procedure_surface.api.base_result import BaseResult
+
+
+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, int | float]
+    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, int | float]
+    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, int | float]
+    configuration: dict[str, Any]
diff --git a/graphdatascience/procedure_surface/arrow/relationship_endpoints_helper.py b/graphdatascience/procedure_surface/arrow/relationship_endpoints_helper.py
@@ -4,7 +4,6 @@
 from graphdatascience.procedure_surface.arrow.endpoints_helper_base import EndpointsHelperBase
 
 
-# TODO find common parts with node_property_endpoints and refactor into a base class
 class RelationshipEndpointsHelper(EndpointsHelperBase):
     """
     Helper class for Arrow algorithm endpoints that work with relationships.
diff --git a/graphdatascience/procedure_surface/arrow/similarity/knn_arrow_endpoints.py b/graphdatascience/procedure_surface/arrow/similarity/knn_arrow_endpoints.py
@@ -8,13 +8,15 @@
 from graphdatascience.arrow_client.v2.remote_write_back_client import RemoteWriteBackClient
 from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
 from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
-from graphdatascience.procedure_surface.api.similarity.knn_endpoints import (
-    KnnEndpoints,
+from graphdatascience.procedure_surface.api.similarity.knn_endpoints import KnnEndpoints
+from graphdatascience.procedure_surface.api.similarity.knn_filtered_endpoints import KnnFilteredEndpoints
+from graphdatascience.procedure_surface.api.similarity.knn_results import (
     KnnMutateResult,
     KnnStatsResult,
     KnnWriteResult,
 )
 from graphdatascience.procedure_surface.arrow.relationship_endpoints_helper import RelationshipEndpointsHelper
+from graphdatascience.procedure_surface.arrow.similarity.knn_filtered_arrow_endpoints import KnnFilteredArrowEndpoints
 from graphdatascience.procedure_surface.arrow.stream_result_mapper import rename_similarity_stream_result
 
 
@@ -29,6 +31,14 @@ def __init__(
             arrow_client, write_back_client=write_back_client, show_progress=show_progress
         )
 
+    @property
+    def filtered(self) -> KnnFilteredEndpoints:
+        return KnnFilteredArrowEndpoints(
+            self._endpoints_helper._arrow_client,
+            self._endpoints_helper._write_back_client,
+            self._endpoints_helper._show_progress,
+        )
+
     def mutate(
         self,
         G: GraphV2,
diff --git a/graphdatascience/procedure_surface/arrow/similarity/knn_filtered_arrow_endpoints.py b/graphdatascience/procedure_surface/arrow/similarity/knn_filtered_arrow_endpoints.py
@@ -6,12 +6,12 @@
 from graphdatascience.arrow_client.v2.remote_write_back_client import RemoteWriteBackClient
 from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
 from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
-from graphdatascience.procedure_surface.api.similarity.knn_endpoints import (
+from graphdatascience.procedure_surface.api.similarity.knn_filtered_endpoints import KnnFilteredEndpoints
+from graphdatascience.procedure_surface.api.similarity.knn_results import (
     KnnMutateResult,
     KnnStatsResult,
     KnnWriteResult,
 )
-from graphdatascience.procedure_surface.api.similarity.knn_filtered_endpoints import KnnFilteredEndpoints
 from graphdatascience.procedure_surface.arrow.relationship_endpoints_helper import RelationshipEndpointsHelper
 from graphdatascience.procedure_surface.arrow.stream_result_mapper import rename_similarity_stream_result
 
diff --git a/graphdatascience/procedure_surface/cypher/similarity/knn_cypher_endpoints.py b/graphdatascience/procedure_surface/cypher/similarity/knn_cypher_endpoints.py
@@ -7,13 +7,17 @@
 from graphdatascience.call_parameters import CallParameters
 from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
 from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
-from graphdatascience.procedure_surface.api.similarity.knn_endpoints import (
-    KnnEndpoints,
+from graphdatascience.procedure_surface.api.similarity.knn_endpoints import KnnEndpoints
+from graphdatascience.procedure_surface.api.similarity.knn_filtered_endpoints import KnnFilteredEndpoints
+from graphdatascience.procedure_surface.api.similarity.knn_results import (
     KnnMutateResult,
     KnnStatsResult,
     KnnWriteResult,
 )
 from graphdatascience.procedure_surface.cypher.estimation_utils import estimate_algorithm
+from graphdatascience.procedure_surface.cypher.similarity.knn_filtered_cypher_endpoints import (
+    KnnFilteredCypherEndpoints,
+)
 from graphdatascience.procedure_surface.utils.config_converter import ConfigConverter
 from graphdatascience.query_runner.query_runner import QueryRunner
 
@@ -22,6 +26,10 @@ class KnnCypherEndpoints(KnnEndpoints):
     def __init__(self, query_runner: QueryRunner):
         self._query_runner = query_runner
 
+    @property
+    def filtered(self) -> KnnFilteredEndpoints:
+        return KnnFilteredCypherEndpoints(self._query_runner)
+
     def mutate(
         self,
         G: GraphV2,
diff --git a/graphdatascience/procedure_surface/cypher/similarity/knn_filtered_cypher_endpoints.py b/graphdatascience/procedure_surface/cypher/similarity/knn_filtered_cypher_endpoints.py
@@ -5,12 +5,12 @@
 from graphdatascience.call_parameters import CallParameters
 from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
 from graphdatascience.procedure_surface.api.estimation_result import EstimationResult
-from graphdatascience.procedure_surface.api.similarity.knn_endpoints import (
+from graphdatascience.procedure_surface.api.similarity.knn_filtered_endpoints import KnnFilteredEndpoints
+from graphdatascience.procedure_surface.api.similarity.knn_results import (
     KnnMutateResult,
     KnnStatsResult,
     KnnWriteResult,
 )
-from graphdatascience.procedure_surface.api.similarity.knn_filtered_endpoints import KnnFilteredEndpoints
 from graphdatascience.procedure_surface.cypher.estimation_utils import estimate_algorithm
 from graphdatascience.procedure_surface.utils.config_converter import ConfigConverter
 from graphdatascience.query_runner.query_runner import QueryRunner
diff --git a/graphdatascience/session/session_v2_endpoints.py b/graphdatascience/session/session_v2_endpoints.py
@@ -16,7 +16,6 @@
 from graphdatascience.procedure_surface.api.community.triangle_count_endpoints import TriangleCountEndpoints
 from graphdatascience.procedure_surface.api.node_embedding.graphsage_endpoints import GraphSageEndpoints
 from graphdatascience.procedure_surface.api.similarity.knn_endpoints import KnnEndpoints
-from graphdatascience.procedure_surface.api.similarity.knn_filtered_endpoints import KnnFilteredEndpoints
 from graphdatascience.procedure_surface.arrow.catalog_arrow_endpoints import CatalogArrowEndpoints
 from graphdatascience.procedure_surface.arrow.centrality.articlerank_arrow_endpoints import ArticleRankArrowEndpoints
 from graphdatascience.procedure_surface.arrow.centrality.articulationpoints_arrow_endpoints import (
@@ -173,10 +172,6 @@ def kmeans(self) -> KMeansEndpoints:
     def knn(self) -> KnnEndpoints:
         return KnnArrowEndpoints(self._arrow_client, self._write_back_client, show_progress=self._show_progress)
 
-    @property
-    def knn_filtered(self) -> KnnFilteredEndpoints:
-        return KnnFilteredArrowEndpoints(self._arrow_client, self._write_back_client, show_progress=self._show_progress)
-
     @property
     def label_propagation(self) -> LabelPropagationEndpoints:
         return LabelPropagationArrowEndpoints(
diff --git a/graphdatascience/tests/integrationV2/procedure_surface/arrow/similarity/test_knn_arrow_endpoints.py b/graphdatascience/tests/integrationV2/procedure_surface/arrow/similarity/test_knn_arrow_endpoints.py
@@ -6,7 +6,6 @@
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
 from graphdatascience.arrow_client.v2.remote_write_back_client import RemoteWriteBackClient
 from graphdatascience.procedure_surface.api.catalog.graph_api import GraphV2
-from graphdatascience.procedure_surface.api.similarity.knn_endpoints import KnnWriteResult
 from graphdatascience.procedure_surface.arrow.similarity.knn_arrow_endpoints import KnnArrowEndpoints
 from graphdatascience.tests.integrationV2.procedure_surface.arrow.graph_creation_helper import (
     create_graph,
@@ -107,7 +106,6 @@ def test_knn_write(arrow_client: AuthenticatedArrowClient, query_runner: QueryRu
         G=db_graph, write_relationship_type="SIMILAR", write_property="similarity", node_properties=["prop"], top_k=2
     )
 
-    assert isinstance(result, KnnWriteResult)
     assert result.ran_iterations > 0
     assert result.did_converge
     assert result.pre_processing_millis >= 0
diff --git a/graphdatascience/tests/integrationV2/procedure_surface/session/test_session_endpoint_coverage.py b/graphdatascience/tests/integrationV2/procedure_surface/session/test_session_endpoint_coverage.py
@@ -57,13 +57,8 @@
 
 
 @pytest.fixture
-def gds(arrow_client: AuthenticatedArrowClient, db_query_runner: QueryRunner) -> AuraGraphDataScience:
-    return AuraGraphDataScience(
-        query_runner=db_query_runner,
-        delete_fn=lambda: True,
-        gds_version=ServerVersion.from_string("2.7.0"),
-        v2_endpoints=SessionV2Endpoints(arrow_client, db_query_runner, show_progress=False),
-    )
+def endpoints(arrow_client: AuthenticatedArrowClient) -> SessionV2Endpoints:
+    return SessionV2Endpoints(arrow_client, db_client=None, show_progress=False)
 
 
 def to_snake(camel: str) -> str:
@@ -98,8 +93,8 @@ def check_gds_v2_availability(endpoints: SessionV2Endpoints, algo: str) -> bool:
 
 
 @pytest.mark.db_integration
-def test_algo_coverage(gds: AuraGraphDataScience) -> None:
-    arrow_client = gds.v2._arrow_client
+def test_algo_coverage(endpoints: SessionV2Endpoints) -> None:
+    arrow_client = endpoints._arrow_client
 
     # Get all available Arrow actions
     available_v2_actions = [
@@ -124,7 +119,7 @@ def test_algo_coverage(gds: AuraGraphDataScience) -> None:
     for category, algos in algos_per_category.items():
         for algo in algos:
             is_available = check_gds_v2_availability(
-                gds.v2,
+                endpoints,
                 algo,
             )
             action = f"{category}.{algo}"
