Add progress bar to fetching FloatSeries and StringSeries (#1633)

Co-authored-by: Rafał Jankowski <rafal.jankowski@neptune.ai>
Co-authored-by: Sabine <sabine.nyholm@neptune.ai>

CHANGELOG.md

@@ -8,6 +8,7 @@
 - Added `ascending` parameter to `fetch_*_table()` methods ([#1602](https://github.com/neptune-ai/neptune-client/pull/1602))
 - Added `progress_bar` parameter to `fetch_*_table()` methods ([#1599](https://github.com/neptune-ai/neptune-client/pull/1599))
 - Added `progress_bar` parameter to `download()` method of the `Handler` class ([#1620](https://github.com/neptune-ai/neptune-client/pull/1620))
+- Added `progress_bar` parameter to `fetch_values()` method of the `Handler` class ([#1633](https://github.com/neptune-ai/neptune-client/pull/1633))
 
 ### Fixes
 - Add direct requirement of `typing-extensions` ([#1586](https://github.com/neptune-ai/neptune-client/pull/1586))

src/neptune/attributes/series/fetchable_series.py

@@ -20,6 +20,7 @@
 from typing import (
     Dict,
     Generic,
+    Optional,
     TypeVar,
     Union,
 )
@@ -28,6 +29,9 @@
     FloatSeriesValues,
     StringSeriesValues,
 )
+from neptune.internal.backends.utils import construct_progress_bar
+from neptune.internal.utils.paths import path_to_str
+from neptune.typing import ProgressBarType
 
 Row = TypeVar("Row", StringSeriesValues, FloatSeriesValues)
 
@@ -37,7 +41,7 @@ class FetchableSeries(Generic[Row]):
     def _fetch_values_from_backend(self, offset, limit) -> Row:
         pass
 
-    def fetch_values(self, *, include_timestamp=True):
+    def fetch_values(self, *, include_timestamp: bool = True, progress_bar: Optional[ProgressBarType] = None):
         import pandas as pd
 
         limit = 1000
@@ -53,10 +57,14 @@ def make_row(entry: Row) -> Dict[str, Union[str, float, datetime]]:
                 row["timestamp"] = datetime.fromtimestamp(entry.timestampMillis / 1000)
             return row
 
-        while offset < val.totalItemCount:
-            batch = self._fetch_values_from_backend(offset, limit)
-            data.extend(batch.values)
-            offset += limit
+        path = path_to_str(self._path) if hasattr(self, "_path") else ""
+        with construct_progress_bar(progress_bar, f"Fetching {path} values") as bar:
+            bar.update(by=len(data), total=val.totalItemCount)  # first fetch before the loop
+            while offset < val.totalItemCount:
+                batch = self._fetch_values_from_backend(offset, limit)
+                data.extend(batch.values)
+                offset += limit
+                bar.update(by=len(batch.values), total=val.totalItemCount)
 
         rows = dict((n, make_row(entry)) for (n, entry) in enumerate(data))
 

src/neptune/handler.py

@@ -557,7 +557,7 @@ def fetch_last(self):
         """
         return self._pass_call_to_attr(function_name="fetch_last")
 
-    def fetch_values(self, *, include_timestamp: Optional[bool] = True):
+    def fetch_values(self, *, include_timestamp: Optional[bool] = True, progress_bar: Optional[ProgressBarType] = None):
         """Fetches all values stored in the series from Neptune.
 
         Available for the following field types:
@@ -568,14 +568,23 @@ def fetch_values(self, *, include_timestamp: Optional[bool] = True):
         Args:
             include_timestamp (bool, optional): Whether the fetched data should include the timestamp field.
                 Defaults to `True`.
+            progress_bar: (bool or Type of progress bar, optional): progress bar to be used while fetching values.
+                If `None` or `True` the default tqdm-based progress bar will be used.
+                If `False` no progress bar will be used.
+                If a type of progress bar is passed, it will be used instead of the default one.
+                Defaults to `None`.
 
         Returns:
             ``Pandas.DataFrame``: containing all the values and their indexes stored in the series field.
 
         For more information on field types, see the docs:
            https://docs.neptune.ai/api-reference/field-types
         """
-        return self._pass_call_to_attr(function_name="fetch_values", include_timestamp=include_timestamp)
+        return self._pass_call_to_attr(
+            function_name="fetch_values",
+            include_timestamp=include_timestamp,
+            progress_bar=progress_bar,
+        )
 
     @check_protected_paths
     def delete_files(self, paths: Union[str, Iterable[str]], *, wait: bool = False) -> None:

src/neptune/internal/backends/hosted_neptune_backend.py

@@ -164,6 +164,15 @@
     AttributeType.RUN_STATE.value,
 }
 
+ATOMIC_ATTRIBUTE_TYPES = {
+    AttributeType.INT.value,
+    AttributeType.FLOAT.value,
+    AttributeType.STRING.value,
+    AttributeType.BOOL.value,
+    AttributeType.DATETIME.value,
+    AttributeType.RUN_STATE.value,
+}
+
 
 class HostedNeptuneBackend(NeptuneBackend):
     def __init__(self, credentials: Credentials, proxies: Optional[Dict[str, str]] = None):