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/arrow_client/v2/gds_arrow_client.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)

graphdatascience/arrow_client/v2/job_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

graphdatascience/procedure_surface/api/catalog/catalog_endpoints.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.

graphdatascience/procedure_surface/arrow/catalog/catalog_arrow_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

graphdatascience/procedure_surface/cypher/catalog_cypher_endpoints.py

@@ -23,6 +23,7 @@
 from graphdatascience.query_runner.arrow_graph_constructor import ArrowGraphConstructor
 from graphdatascience.query_runner.cypher_graph_constructor import CypherGraphConstructor
 from graphdatascience.query_runner.graph_constructor import GraphConstructor
+from graphdatascience.query_runner.neo4j_query_runner import Neo4jQueryRunner
 
 from ...call_parameters import CallParameters
 from ..api.base_result import BaseResult
@@ -34,9 +35,8 @@
 
 
 class CatalogCypherEndpoints(CatalogEndpoints):
-    def __init__(self, cypher_runner: Neo4jQueryRunner, arrow_client: GdsArrowClient | None = None, arrow_client: GdsArrowClient | None = None):
-        self.cypher_runner = cypher_runner
-        self._arrow_client = arrow_client
+    def __init__(self, cypher_runner: Neo4jQueryRunner, arrow_client: GdsArrowClient | None = None):
+        self._cypher_runner = cypher_runner
         self._arrow_client = arrow_client
 
     def construct(
@@ -56,7 +56,7 @@ def construct(
 
         graph_constructor: GraphConstructor
         if self._arrow_client is not None:
-            database = require_database(self._query_runner)
+            database = require_database(self._cypher_runner)
 
             graph_constructor = ArrowGraphConstructor(
                 database=database,
@@ -67,20 +67,20 @@ def construct(
             )
         else:
             graph_constructor = CypherGraphConstructor(
-                query_runner=self._query_runner,
+                query_runner=self._cypher_runner,
                 graph_name=graph_name,
                 concurrency=concurrency,
                 undirected_relationship_types=undirected_relationship_types,
             )
 
         graph_constructor.run(node_dfs=nodes, relationship_dfs=relationships)
-        return GraphV2(name=graph_name, backend=CypherGraphBackend(graph_name, self._query_runner))
+        return GraphV2(name=graph_name, backend=CypherGraphBackend(graph_name, self._cypher_runner))
 
     def list(self, G: GraphV2 | str | None = None) -> list[GraphInfoWithDegrees]:
         graph_name = G if isinstance(G, str) else G.name() if G is not None else None
         params = CallParameters(graphName=graph_name) if graph_name else CallParameters()
 
-        result = self.cypher_runner.call_procedure(endpoint="gds.graph.list", params=params)
+        result = self._cypher_runner.call_procedure(endpoint="gds.graph.list", params=params)
         return [GraphInfoWithDegrees(**row.to_dict()) for _, row in result.iterrows()]
 
     def drop(self, G: GraphV2 | str, fail_if_missing: bool = True) -> GraphInfo | None:
@@ -92,7 +92,7 @@ def drop(self, G: GraphV2 | str, fail_if_missing: bool = True) -> GraphInfo | No
             else CallParameters(graphName=graph_name)
         )
 
-        result = self.cypher_runner.call_procedure(endpoint="gds.graph.drop", params=params)
+        result = self._cypher_runner.call_procedure(endpoint="gds.graph.drop", params=params)
         if len(result) > 0:
             return GraphInfo(**result.iloc[0].to_dict())
         else:
@@ -128,11 +128,11 @@ def project(
         )
         params.ensure_job_id_in_config()
 
-        result = self.cypher_runner.call_procedure(
+        result = self._cypher_runner.call_procedure(
             endpoint="gds.graph.project", params=params, logging=log_progress
         ).squeeze()
         project_result = GraphProjectResult(**result.to_dict())
-        return GraphWithProjectResult(get_graph(project_result.graph_name, self.cypher_runner), project_result)
+        return GraphWithProjectResult(get_graph(project_result.graph_name, self._cypher_runner), project_result)
 
     def filter(
         self,
@@ -158,10 +158,10 @@ def filter(
         )
         params.ensure_job_id_in_config()
 
-        result = self.cypher_runner.call_procedure(
+        result = self._cypher_runner.call_procedure(
             endpoint="gds.graph.filter", params=params, logging=log_progress
         ).squeeze()
-        return GraphWithFilterResult(get_graph(graph_name, self.cypher_runner), GraphFilterResult(**result.to_dict()))
+        return GraphWithFilterResult(get_graph(graph_name, self._cypher_runner), GraphFilterResult(**result.to_dict()))
 
     def generate(
         self,
@@ -202,28 +202,28 @@ def generate(
 
         params.ensure_job_id_in_config()
 
-        result = self.cypher_runner.call_procedure(
+        result = self._cypher_runner.call_procedure(
             endpoint="gds.graph.generate", params=params, logging=log_progress
         ).squeeze()
         return GraphWithGenerationStats(
-            get_graph(graph_name, self.cypher_runner), GraphGenerationStats(**result.to_dict())
+            get_graph(graph_name, self._cypher_runner), GraphGenerationStats(**result.to_dict())
         )
 
     @property
     def sample(self) -> GraphSamplingEndpoints:
-        return GraphSamplingCypherEndpoints(self.cypher_runner)
+        return GraphSamplingCypherEndpoints(self._cypher_runner)
 
     @property
     def node_labels(self) -> NodeLabelCypherEndpoints:
-        return NodeLabelCypherEndpoints(self.cypher_runner)
+        return NodeLabelCypherEndpoints(self._cypher_runner)
 
     @property
     def node_properties(self) -> NodePropertiesCypherEndpoints:
-        return NodePropertiesCypherEndpoints(self.cypher_runner, self._arrow_client)
+        return NodePropertiesCypherEndpoints(self._cypher_runner, self._arrow_client)
 
     @property
     def relationships(self) -> RelationshipCypherEndpoints:
-        return RelationshipCypherEndpoints(self.cypher_runner, self._arrow_client)
+        return RelationshipCypherEndpoints(self._cypher_runner, self._arrow_client)
 
 
 class GraphProjectResult(BaseResult):