Implement graph_construct using arrow v2

* bringing terminationFlag to GdsArrowClient (V2) to interrupt upload
* job client also support waiting for a given status

Author: FlorentinD

graphdatascience/tests/unit/procedure_surface/arrow/__init__.py

@@ -0,0 +1,120 @@
+from contextlib import ExitStack
+from unittest import mock
+
+from pandas import DataFrame
+from pytest_mock import MockerFixture
+
+from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from graphdatascience.arrow_client.v2.api_types import JobStatus
+from graphdatascience.procedure_surface.arrow.catalog.catalog_arrow_endpoints import CatalogArrowEndpoints
+from graphdatascience.tests.unit.arrow_client.arrow_test_utils import ArrowTestResult
+
+
+def test_construct_with_no_rels(mocker: MockerFixture) -> None:
+    arrow_client = mocker.Mock(spec=AuthenticatedArrowClient)
+    job_id = "job-123"
+
+    relationship_loading_done_status = JobStatus(
+        jobId=job_id,
+        status="RELATIONSHIP_LOADING",
+        progress=-1,
+        description="",
+    )
+    construct_done_status = JobStatus(
+        jobId=job_id,
+        status="Done",
+        progress=-1,
+        description="",
+    )
+
+    do_action_with_retry = mocker.Mock()
+    do_action_with_retry.side_effect = [
+        iter([ArrowTestResult(relationship_loading_done_status.dump_camel())]),
+        iter([ArrowTestResult(construct_done_status.dump_camel())]),
+    ]
+
+    arrow_client.do_action_with_retry = do_action_with_retry
+
+    endpoints = CatalogArrowEndpoints(arrow_client=arrow_client)
+
+    nodes = DataFrame(
+        {
+            "nodeId": [0, 1],
+            "labels": [["A"], ["B"]],
+            "propA": [1337, 42.1],
+        }
+    )
+    with patch_gds_arrow_client(job_id):
+        G = endpoints.construct(graph_name="g", nodes=nodes, relationships=[])
+        assert G.name() == "g"
+
+
+def test_construct_with_df_lists(mocker: MockerFixture) -> None:
+    arrow_client = mocker.Mock(spec=AuthenticatedArrowClient)
+    job_id = "foo"
+    relationship_loading_done_status = JobStatus(
+        jobId=job_id,
+        status="RELATIONSHIP_LOADING",
+        progress=-1,
+        description="",
+    )
+    construct_done_status = JobStatus(
+        jobId=job_id,
+        status="Done",
+        progress=-1,
+        description="",
+    )
+
+    do_action_with_retry = mocker.Mock()
+    do_action_with_retry.side_effect = [
+        iter([ArrowTestResult(relationship_loading_done_status.dump_camel())]),
+        iter([ArrowTestResult(construct_done_status.dump_camel())]),
+    ]
+    arrow_client.do_action_with_retry = do_action_with_retry
+
+    endpoints = CatalogArrowEndpoints(arrow_client=arrow_client)
+
+    nodes = [
+        DataFrame({"nodeId": [0, 1], "labels": ["a", "a"], "property": [6.0, 7.0]}),
+        DataFrame({"nodeId": [2, 3], "labels": ["b", "b"], "q": [-500, -400]}),
+    ]
+    relationships = [
+        DataFrame(
+            {"sourceNodeId": [0, 1], "targetNodeId": [1, 2], "relationshipType": ["A", "A"], "weights": [0.2, 0.3]}
+        ),
+        DataFrame({"sourceNodeId": [2, 3], "targetNodeId": [3, 0], "relationshipType": ["B", "B"]}),
+    ]
+    with patch_gds_arrow_client(job_id):
+        G = endpoints.construct(graph_name="g", nodes=nodes, relationships=relationships)
+        assert G.name() == "g"
+
+
+def patch_gds_arrow_client(create_graph_job_id: str) -> ExitStack:
+    exit_stack = ExitStack()
+    patches = [
+        mock.patch(
+            "graphdatascience.arrow_client.v2.gds_arrow_client.GdsArrowClient.create_graph",
+            return_value=create_graph_job_id,
+        ),
+        mock.patch(
+            "graphdatascience.arrow_client.v2.gds_arrow_client.GdsArrowClient.upload_nodes",
+            return_value=None,
+        ),
+        mock.patch(
+            "graphdatascience.arrow_client.v2.gds_arrow_client.GdsArrowClient.upload_relationships",
+            return_value=None,
+        ),
+        mock.patch(
+            "graphdatascience.arrow_client.v2.gds_arrow_client.GdsArrowClient.node_load_done",
+            return_value=None,
+        ),
+        mock.patch(
+            "graphdatascience.arrow_client.v2.gds_arrow_client.GdsArrowClient.relationship_load_done",
+            return_value=None,
+        ),
+    ]
+
+    for p in patches:
+        exit_stack.enter_context(p)
+
+    return exit_stack

graphdatascience/tests/unit/procedure_surface/arrow/test_catalog_arrow_endpoints.py

@@ -11,6 +11,7 @@
 
 from graphdatascience.arrow_client.arrow_endpoint_version import ArrowEndpointVersion
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient, ConnectionInfo
+from graphdatascience.query_runner.termination_flag import TerminationFlag
 
 from ...procedure_surface.api.default_values import ALL_TYPES
 from ...procedure_surface.utils.config_converter import ConfigConverter
@@ -328,6 +329,7 @@ def upload_nodes(
         data: pyarrow.Table | list[pyarrow.RecordBatch] | pandas.DataFrame,
         batch_size: int = 10000,
         progress_callback: Callable[[int], None] = lambda x: None,
+        termination_flag: TerminationFlag | None = None,
     ) -> None:
         """
         Uploads node data to the server for a given job.
@@ -342,15 +344,20 @@ def upload_nodes(
             The number of rows per batch
         progress_callback
             A callback function that is called with the number of rows uploaded after each batch
+        termination_flag
+            A termination flag to cancel the upload if requested
         """
-        self._upload_data("graph.project.fromTables.nodes", job_id, data, batch_size, progress_callback)
+        self._upload_data(
+            "graph.project.fromTables.nodes", job_id, data, batch_size, progress_callback, termination_flag
+        )
 
     def upload_relationships(
         self,
         job_id: str,
         data: pyarrow.Table | list[pyarrow.RecordBatch] | pandas.DataFrame,
         batch_size: int = 10000,
         progress_callback: Callable[[int], None] = lambda x: None,
+        termination_flag: TerminationFlag | None = None,
     ) -> None:
         """
         Uploads relationship data to the server for a given job.
@@ -365,15 +372,20 @@ def upload_relationships(
             The number of rows per batch
         progress_callback
             A callback function that is called with the number of rows uploaded after each batch
+        termination_flag
+            A termination flag to cancel the upload if requested
         """
-        self._upload_data("graph.project.fromTables.relationships", job_id, data, batch_size, progress_callback)
+        self._upload_data(
+            "graph.project.fromTables.relationships", job_id, data, batch_size, progress_callback, termination_flag
+        )
 
     def upload_triplets(
         self,
         job_id: str,
         data: pyarrow.Table | list[pyarrow.RecordBatch] | pandas.DataFrame,
         batch_size: int = 10000,
         progress_callback: Callable[[int], None] = lambda x: None,
+        termination_flag: TerminationFlag | None = None,
     ) -> None:
         """
         Uploads triplet data to the server for a given job.
@@ -388,8 +400,10 @@ def upload_triplets(
             The number of rows per batch
         progress_callback
             A callback function that is called with the number of rows uploaded after each batch
+        termination_flag
+            A termination flag to cancel the upload if requested
         """
-        self._upload_data("graph.project.fromTriplets", job_id, data, batch_size, progress_callback)
+        self._upload_data("graph.project.fromTriplets", job_id, data, batch_size, progress_callback, termination_flag)
 
     def abort_job(self, job_id: str) -> None:
         """
@@ -464,6 +478,7 @@ def _upload_data(
         data: pyarrow.Table | list[pyarrow.RecordBatch] | pandas.DataFrame,
         batch_size: int = 10000,
         progress_callback: Callable[[int], None] = lambda x: None,
+        termination_flag: TerminationFlag | None = None,
     ) -> None:
         match data:
             case pyarrow.Table():
@@ -490,6 +505,10 @@ def upload_batch(p: RecordBatch) -> None:
 
         with put_stream:
             for partition in batches:
+                if termination_flag is not None and termination_flag.is_set():
+                    self.abort_job(job_id)  # closing the put_stream will raise an error
+                    break
+
                 upload_batch(partition)
                 ack_stream.read()
                 progress_callback(partition.num_rows)

src/graphdatascience/arrow_client/v2/gds_arrow_client.py

@@ -39,20 +39,26 @@ def wait_for_job(
         client: AuthenticatedArrowClient,
         job_id: str,
         show_progress: bool,
+        expected_status: str | None = None,
         termination_flag: TerminationFlag | None = None,
     ) -> None:
         progress_bar: TqdmProgressBar | None = None
 
+        def check_expected_status(status: JobStatus) -> bool:
+            return job_status.succeeded() if expected_status is None else status.status == expected_status
+
         if termination_flag is None:
             termination_flag = TerminationFlag.create()
 
-        for attempt in Retrying(retry=retry_if_result(lambda _: True), wait=wait_exponential(min=0.1, max=5)):
+        for attempt in Retrying(
+            retry=retry_if_result(lambda _: True), wait=wait_exponential(min=0.1, max=5), reraise=True
+        ):
             with attempt:
                 termination_flag.assert_running()
 
                 job_status = self.get_job_status(client, job_id)
 
-                if job_status.succeeded() or job_status.aborted():
+                if check_expected_status(job_status) or job_status.aborted():
                     if progress_bar:
                         progress_bar.finish(success=job_status.succeeded())
                     return

src/graphdatascience/arrow_client/v2/job_client.py

@@ -22,7 +22,7 @@ def construct(
         graph_name: str,
         nodes: DataFrame | list[DataFrame],
         relationships: DataFrame | list[DataFrame] | None = None,
-        concurrency: int = 4,
+        concurrency: int | None = None,
         undirected_relationship_types: list[str] | None = None,
     ) -> GraphV2:
         """Construct a graph from a list of node and relationship dataframes.

src/graphdatascience/procedure_surface/api/catalog/catalog_endpoints.py

@@ -8,6 +8,7 @@
 from pandas import DataFrame
 
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
+from graphdatascience.arrow_client.v2.gds_arrow_client import GdsArrowClient
 from graphdatascience.arrow_client.v2.job_client import JobClient
 from graphdatascience.arrow_client.v2.remote_write_back_client import RemoteWriteBackClient
 from graphdatascience.procedure_surface.api.base_result import BaseResult
@@ -31,6 +32,7 @@
 )
 from graphdatascience.procedure_surface.arrow.catalog.relationship_arrow_endpoints import RelationshipArrowEndpoints
 from graphdatascience.procedure_surface.utils.config_converter import ConfigConverter
+from graphdatascience.query_runner.progress.progress_bar import NoOpProgressBar, ProgressBar, TqdmProgressBar
 from graphdatascience.query_runner.protocol.project_protocols import ProjectProtocol
 from graphdatascience.query_runner.query_runner import QueryRunner
 from graphdatascience.query_runner.termination_flag import TerminationFlag
@@ -135,10 +137,73 @@ def construct(
         graph_name: str,
         nodes: DataFrame | list[DataFrame],
         relationships: DataFrame | list[DataFrame] | None = None,
-        concurrency: int = 4,
+        concurrency: int | None = None,
         undirected_relationship_types: list[str] | None = None,
     ) -> GraphV2:
-        raise NotImplementedError("Graph construction is not yet supported via V2 endpoints.")
+        gds_arrow_client = GdsArrowClient(self._arrow_client)
+        job_client = JobClient()
+        termination_flag = TerminationFlag.create()
+
+        if self._show_progress:
+            progress_bar: ProgressBar = TqdmProgressBar(task_name="Constructing graph", relative_progress=0.0)
+        else:
+            progress_bar = NoOpProgressBar()
+
+        with progress_bar:
+            create_job_id: str = gds_arrow_client.create_graph(
+                graph_name=graph_name,
+                undirected_relationship_types=undirected_relationship_types or [],
+                concurrency=concurrency,
+            )
+            node_count = nodes.shape[0] if isinstance(nodes, DataFrame) else sum(df.shape[0] for df in nodes)
+            if isinstance(relationships, DataFrame):
+                rel_count = relationships.shape[0]
+            elif relationships is None:
+                rel_count = 0
+                relationships = []
+            else:
+                rel_count = sum(df.shape[0] for df in relationships)
+            total_count = node_count + rel_count
+
+            gds_arrow_client.upload_nodes(
+                create_job_id,
+                nodes,
+                progress_callback=lambda rows_imported: progress_bar.update(
+                    sub_tasks_description="Uploading nodes", progress=rows_imported / total_count, status="Running"
+                ),
+                termination_flag=termination_flag,
+            )
+
+            gds_arrow_client.node_load_done(create_job_id)
+
+            # skipping progress bar here as we have our own for the overall process
+            job_client.wait_for_job(
+                self._arrow_client,
+                create_job_id,
+                expected_status="RELATIONSHIP_LOADING",
+                termination_flag=termination_flag,
+                show_progress=False,
+            )
+
+            if rel_count > 0:
+                gds_arrow_client.upload_relationships(
+                    create_job_id,
+                    relationships,
+                    progress_callback=lambda rows_imported: progress_bar.update(
+                        sub_tasks_description="Uploading relationships",
+                        progress=rows_imported / total_count,
+                        status="Running",
+                    ),
+                    termination_flag=termination_flag,
+                )
+
+            gds_arrow_client.relationship_load_done(create_job_id)
+
+        # will produce a second progress bar to show graph construction on the server side
+        job_client.wait_for_job(
+            self._arrow_client, create_job_id, termination_flag=termination_flag, show_progress=True
+        )
+        return get_graph(graph_name, self._arrow_client)
 
     def drop(self, G: GraphV2 | str, fail_if_missing: bool = True) -> GraphInfo | None:
         graph_name = G.name() if isinstance(G, GraphV2) else G