Fill all progress bar on completion

Author: FlorentinD

changelog.md

@@ -20,6 +20,7 @@
 * Fixed an issue where source and target IDs of relationships in heterogeneous OGBL graphs were not parsed correctly.
 * Fixed an issue where configuration parameters such as `aggregation` were ignored by `gds.graph.toUndirected`.
 * Fixed an issue where the `database` given for the `GraphDataScience` construction was not used for metadata retrieval, causing an exception to be raised if the default "neo4j" database was missing.
+* Fixed an issue where progress bars would not complete.
 
 
 ## Improvements

graphdatascience/query_runner/arrow_graph_constructor.py

@@ -3,7 +3,7 @@
 import math
 import warnings
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, NoReturn, Optional
 
 import numpy
 import pyarrow.flight as flight
@@ -93,7 +93,7 @@ def _send_action(self, action_type: str, meta_data: Dict[str, Any]) -> None:
 
         json.loads(collected_result[0].body.to_pybytes().decode())
 
-    def _send_df(self, df: DataFrame, entity_type: str, pbar: tqdm) -> None:
+    def _send_df(self, df: DataFrame, entity_type: str, pbar: tqdm[NoReturn]) -> None:
         table = Table.from_pandas(df)
         batches = table.to_batches(self._chunk_size)
         flight_descriptor = {"name": self._graph_name, "entity_type": entity_type}

graphdatascience/query_runner/neo4j_query_runner.py

@@ -5,7 +5,7 @@
 import time
 import warnings
 from concurrent.futures import Future, ThreadPoolExecutor, wait
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, NoReturn, Optional, Tuple, Union
 from uuid import uuid4
 
 import neo4j
@@ -226,7 +226,7 @@ def _forward_cypher_warnings(self, notification: Dict[str, Any]) -> None:
             self._logger.info(notification)
 
     def _log(self, job_id: str, future: "Future[Any]", database: Optional[str] = None) -> None:
-        pbar = None
+        pbars: Dict[str, tqdm[NoReturn]] = {}
         warn_if_failure = True
 
         while wait([future], timeout=self._LOG_POLLING_INTERVAL).not_done:
@@ -248,16 +248,18 @@ def _log(self, job_id: str, future: "Future[Any]", database: Optional[str] = Non
                     continue
 
             progress_percent = progress["progress"][0]
-            if not progress_percent == "n/a":
-                task_name = progress["taskName"][0].split("|--")[-1][1:]
-                pbar = pbar or tqdm(total=100, unit="%", desc=task_name)
-            else:
+            if progress_percent == "n/a":
                 return
 
+            task_name = progress["taskName"][0].split("|--")[-1][1:]
+            if task_name not in pbars:
+                pbars[task_name] = tqdm(total=100, unit="%", desc=task_name)
+            pbar = pbars[task_name]
+
             parsed_progress = float(progress_percent[:-1])
             pbar.update(parsed_progress - pbar.n)
 
-        if pbar:
+        for pbar in pbars.values():
             pbar.update(100 - pbar.n)
 
     def set_database(self, database: str) -> None:

requirements/dev/dev.txt

@@ -10,3 +10,4 @@ tox == 4.11.3
 types-setuptools == 68.1.0.1
 sphinx == 7.2.6
 types-requests
+types-tqdm