Support stream mode

with the latest session release we can call the endpoints without a seg fault

Author: FlorentinD

graphdatascience/procedure_surface/arrow/similarity/knn_arrow_endpoints.py

@@ -15,6 +15,7 @@
     KnnWriteResult,
 )
 from graphdatascience.procedure_surface.arrow.relationship_endpoints_helper import RelationshipEndpointsHelper
+from graphdatascience.procedure_surface.arrow.stream_result_mapper import rename_similarity_stream_result
 
 
 class KnnArrowEndpoints(KnnEndpoints):
@@ -145,29 +146,30 @@ def stream(
         concurrency: Any | None = None,
         job_id: Any | None = None,
     ) -> DataFrame:
-        # config = self._endpoints_helper.create_base_config(
-        #     G,
-        #     nodeProperties=node_properties,
-        #     topK=top_k,
-        #     similarityCutoff=similarity_cutoff,
-        #     deltaThreshold=delta_threshold,
-        #     maxIterations=max_iterations,
-        #     sampleRate=sample_rate,
-        #     perturbationRate=perturbation_rate,
-        #     randomJoins=random_joins,
-        #     randomSeed=random_seed,
-        #     initialSampler=initial_sampler,
-        #     relationshipTypes=relationship_types,
-        #     nodeLabels=node_labels,
-        #     sudo=sudo,
-        #     logProgress=log_progress,
-        #     username=username,
-        #     concurrency=concurrency,
-        #     jobId=job_id,
-        # )
-        # return self._endpoints_helper.run_job_and_stream("v2/similarity.knn", G, config)
+        config = self._endpoints_helper.create_base_config(
+            G,
+            nodeProperties=node_properties,
+            topK=top_k,
+            similarityCutoff=similarity_cutoff,
+            deltaThreshold=delta_threshold,
+            maxIterations=max_iterations,
+            sampleRate=sample_rate,
+            perturbationRate=perturbation_rate,
+            randomJoins=random_joins,
+            randomSeed=random_seed,
+            initialSampler=initial_sampler,
+            relationshipTypes=relationship_types,
+            nodeLabels=node_labels,
+            sudo=sudo,
+            logProgress=log_progress,
+            username=username,
+            concurrency=concurrency,
+            jobId=job_id,
+        )
+        result = self._endpoints_helper.run_job_and_stream("v2/similarity.knn", G, config)
+        rename_similarity_stream_result(result)
 
-        raise NotImplementedError()
+        return result
 
     def write(
         self,

graphdatascience/procedure_surface/arrow/similarity/knn_filtered_arrow_endpoints.py

@@ -13,6 +13,7 @@
 )
 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
 
 
 class KnnFilteredArrowEndpoints(KnnFilteredEndpoints):
@@ -154,7 +155,34 @@ def stream(
         concurrency: Any | None = None,
         job_id: Any | None = None,
     ) -> DataFrame:
-        raise NotImplementedError("Filtered KNN stream endpoint is not available via Arrow")
+        config = self._endpoints_helper.create_base_config(
+            G,
+            nodeProperties=node_properties,
+            sourceNodeFilter=source_node_filter,
+            targetNodeFilter=target_node_filter,
+            seedTargetNodes=seed_target_nodes,
+            nodeLabels=node_labels,
+            relationshipTypes=relationship_types,
+            similarityCutoff=similarity_cutoff,
+            perturbationRate=perturbation_rate,
+            deltaThreshold=delta_threshold,
+            sampleRate=sample_rate,
+            randomJoins=random_joins,
+            initialSampler=initial_sampler,
+            maxIterations=max_iterations,
+            topK=top_k,
+            randomSeed=random_seed,
+            concurrency=concurrency,
+            jobId=job_id,
+            logProgress=log_progress,
+            sudo=sudo,
+            username=username,
+        )
+
+        result = self._endpoints_helper.run_job_and_stream("v2/similarity.knn.filtered", G, config)
+        rename_similarity_stream_result(result)
+
+        return result
 
     def write(
         self,

graphdatascience/procedure_surface/arrow/stream_result_mapper.py

@@ -0,0 +1,7 @@
+from pandas import DataFrame
+
+
+def rename_similarity_stream_result(result: DataFrame) -> None:
+    result.rename(columns={"sourceNodeId": "node1", "targetNodeId": "node2"}, inplace=True)
+    if "relationshipType" in result.columns:
+        result.drop(columns=["relationshipType"], inplace=True)

graphdatascience/tests/integrationV2/procedure_surface/arrow/similarity/test_knn_arrow_endpoints.py

@@ -64,7 +64,6 @@ def test_knn_stats(knn_endpoints: KnnArrowEndpoints, sample_graph: GraphV2) -> N
     assert "p50" in result.similarity_distribution
 
 
-@pytest.mark.skip(reason="SEGFAULT for custom metadata. tracked in GDSA-312")
 def test_knn_stream(knn_endpoints: KnnArrowEndpoints, sample_graph: GraphV2) -> None:
     """Test KNN stream operation."""
     result_df = knn_endpoints.stream(
@@ -74,7 +73,7 @@ def test_knn_stream(knn_endpoints: KnnArrowEndpoints, sample_graph: GraphV2) ->
     )
 
     assert set(result_df.columns) == {"node1", "node2", "similarity"}
-    assert len(result_df) == 2
+    assert len(result_df) == 8
 
 
 def test_knn_mutate(knn_endpoints: KnnArrowEndpoints, sample_graph: GraphV2) -> None:

graphdatascience/tests/integrationV2/procedure_surface/arrow/similarity/test_knn_filtered_arrow_endpoints.py

@@ -49,7 +49,7 @@ def knn_filtered_endpoints(arrow_client: AuthenticatedArrowClient) -> Generator[
     yield KnnFilteredArrowEndpoints(arrow_client)
 
 
-def test_stats(knn_filtered_endpoints: KnnFilteredArrowEndpoints, sample_graph: GraphV2) -> None:
+def test_knn_filtered_stats(knn_filtered_endpoints: KnnFilteredArrowEndpoints, sample_graph: GraphV2) -> None:
     result = knn_filtered_endpoints.stats(
         sample_graph,
         node_properties="prop",
@@ -70,10 +70,7 @@ def test_stats(knn_filtered_endpoints: KnnFilteredArrowEndpoints, sample_graph:
     assert result.configuration is not None
 
 
-@pytest.mark.skip(reason="SEGFAULT for custom metadata. tracked in GDSA-312")
-def test_stream_raises_not_implemented(
-    knn_filtered_endpoints: KnnFilteredArrowEndpoints, sample_graph: GraphV2
-) -> None:
+def test_knn_filtered_stream(knn_filtered_endpoints: KnnFilteredArrowEndpoints, sample_graph: GraphV2) -> None:
     result_df = knn_filtered_endpoints.stream(
         G=sample_graph,
         node_properties=["prop"],
@@ -83,10 +80,10 @@ def test_stream_raises_not_implemented(
     )
 
     assert set(result_df.columns) == {"node1", "node2", "similarity"}
-    assert len(result_df) == 2
+    assert len(result_df) == 4
 
 
-def test_mutate(knn_filtered_endpoints: KnnFilteredArrowEndpoints, sample_graph: GraphV2) -> None:
+def test_knn_filtered_mutate(knn_filtered_endpoints: KnnFilteredArrowEndpoints, sample_graph: GraphV2) -> None:
     result = knn_filtered_endpoints.mutate(
         sample_graph,
         node_properties="prop",
@@ -110,7 +107,9 @@ def test_mutate(knn_filtered_endpoints: KnnFilteredArrowEndpoints, sample_graph:
     assert result.configuration is not None
 
 
-def test_knn_write(arrow_client: AuthenticatedArrowClient, query_runner: QueryRunner, db_graph: GraphV2) -> None:
+def test_knn_filtered_write(
+    arrow_client: AuthenticatedArrowClient, query_runner: QueryRunner, db_graph: GraphV2
+) -> None:
     endpoints = KnnFilteredArrowEndpoints(
         arrow_client, write_back_client=RemoteWriteBackClient(arrow_client, query_runner), show_progress=False
     )
@@ -139,7 +138,7 @@ def test_knn_write(arrow_client: AuthenticatedArrowClient, query_runner: QueryRu
     assert result.configuration is not None
 
 
-def test_estimate(knn_filtered_endpoints: KnnFilteredArrowEndpoints, sample_graph: GraphV2) -> None:
+def test_knn_filtered_estimate(knn_filtered_endpoints: KnnFilteredArrowEndpoints, sample_graph: GraphV2) -> None:
     result = knn_filtered_endpoints.estimate(
         sample_graph,
         node_properties="prop",