feat: add progress bar for to_arrow method

google/cloud/bigquery/_tqdm_helpers.py

@@ -0,0 +1,49 @@
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared helper functions for tqdm progress bar."""
+
+import warnings
+
+try:
+    import tqdm
+except ImportError:  # pragma: NO COVER
+    tqdm = None
+
+_NO_TQDM_ERROR = (
+    "A progress bar was requested, but there was an error loading the tqdm "
+    "library. Please install tqdm to use the progress bar functionality."
+)
+
+
+def _get_progress_bar(progress_bar_type, description, total, unit):
+    """Construct a tqdm progress bar object, if tqdm is installed."""
+    if tqdm is None:
+        if progress_bar_type is not None:
+            warnings.warn(_NO_TQDM_ERROR, UserWarning, stacklevel=3)
+        return None
+
+    try:
+        if progress_bar_type == "tqdm":
+            return tqdm.tqdm(desc=description, total=total, unit=unit)
+        elif progress_bar_type == "tqdm_notebook":
+            return tqdm.tqdm_notebook(desc=description, total=total, unit=unit)
+        elif progress_bar_type == "tqdm_gui":
+            return tqdm.tqdm_gui(desc=description, total=total, unit=unit)
+    except (KeyError, TypeError):
+        # Protect ourselves from any tqdm errors. In case of
+        # unexpected tqdm behavior, just fall back to showing
+        # no progress bar.
+        warnings.warn(_NO_TQDM_ERROR, UserWarning, stacklevel=3)
+    return None

google/cloud/bigquery/job.py

@@ -20,6 +20,7 @@
 import copy
 import re
 import threading
+import time
 
 import requests
 import six
@@ -52,6 +53,7 @@
 from google.cloud.bigquery.table import TableListItem
 from google.cloud.bigquery.table import TableReference
 from google.cloud.bigquery.table import TimePartitioning
+from google.cloud.bigquery._tqdm_helpers import _get_progress_bar
 
 _DONE_STATE = "DONE"
 _STOPPED_REASON = "stopped"
@@ -3335,7 +3337,53 @@ def to_arrow(
 
         ..versionadded:: 1.17.0
         """
-        return self.result().to_arrow(
+        if progress_bar_type:
+            start_time = time.time()
+            progress_bar = _get_progress_bar(
+                progress_bar_type, "Query is running", 1, "query"
+            )
+            if self.query_plan:
+                i = 0
+                while True:
+                    total = len(self.query_plan)
+                    self.reload()  # Refreshes the state via a GET request.
+                    current_stage = self.query_plan[i]
+                    progress_bar.total = len(self.query_plan)
+                    progress_bar.set_description(
+                        "Query executing stage {} and status {} : {:0.2f}s".format(
+                            current_stage.name,
+                            current_stage.status,
+                            time.time() - start_time,
+                        ),
+                    )
+
+                    try:
+                        query_result = self.result(timeout=0.5)
+                        progress_bar.update(total)
+                        progress_bar.set_description(
+                            "Query complete after {:0.2f}s".format(
+                                time.time() - start_time
+                            ),
+                        )
+                        break
+                    except concurrent.futures.TimeoutError:
+                        if current_stage.status == "COMPLETE":
+                            if i < total - 1:
+                                progress_bar.update(i + 1)
+                                i += 1
+                        continue
+
+            else:
+                query_result = self.result()
+                progress_bar.set_description(
+                    "Query complete after {:0.2f}s".format(time.time() - start_time),
+                )
+                progress_bar.update(1)
+            progress_bar.close()
+        else:
+            query_result = self.result()
+
+        return query_result.to_arrow(
             progress_bar_type=progress_bar_type,
             bqstorage_client=bqstorage_client,
             create_bqstorage_client=create_bqstorage_client,
@@ -3404,7 +3452,51 @@ def to_dataframe(
         Raises:
             ValueError: If the `pandas` library cannot be imported.
         """
-        return self.result().to_dataframe(
+        if progress_bar_type:
+            start_time = time.time()
+            progress_bar = _get_progress_bar(
+                progress_bar_type, "Query is running", 1, "query"
+            )
+            if self.query_plan:
+                i = 0
+                while True:
+                    total = len(self.query_plan)
+                    self.reload()  # Refreshes the state via a GET request.
+
+                    current_stage = self.query_plan[i]
+                    progress_bar.set_description(
+                        "Query executing stage {} and status {} : {:0.2f}s".format(
+                            current_stage.name,
+                            current_stage.status,
+                            time.time() - start_time,
+                        )
+                    )
+
+                    try:
+                        query_result = self.result(timeout=0.5)
+                        progress_bar.update(total)
+                        progress_bar.set_description(
+                            "Query complete after {:0.2f}s".format(
+                                time.time() - start_time
+                            )
+                        )
+                        break
+                    except concurrent.futures.TimeoutError:
+                        if current_stage.status == "COMPLETE":
+                            if i < total - 1:
+                                progress_bar.update(i + 1)
+                                i += 1
+                        continue
+            else:
+                query_result = self.result()
+                progress_bar.set_description(
+                    "Query complete after {:0.2f}s".format(time.time() - start_time),
+                )
+                progress_bar.update(1)
+            progress_bar.close()
+        else:
+            query_result = self.result()
+        return query_result.to_dataframe(
             bqstorage_client=bqstorage_client,
             dtypes=dtypes,
             progress_bar_type=progress_bar_type,

google/cloud/bigquery/table.py

@@ -36,11 +36,6 @@
 except ImportError:  # pragma: NO COVER
     pyarrow = None
 
-try:
-    import tqdm
-except ImportError:  # pragma: NO COVER
-    tqdm = None
-
 import google.api_core.exceptions
 from google.api_core.page_iterator import HTTPIterator
 
@@ -50,6 +45,7 @@
 from google.cloud.bigquery.schema import _build_schema_resource
 from google.cloud.bigquery.schema import _parse_schema_resource
 from google.cloud.bigquery.schema import _to_schema_fields
+from google.cloud.bigquery._tqdm_helpers import _get_progress_bar
 from google.cloud.bigquery.external_config import ExternalConfig
 from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration
 
@@ -68,10 +64,7 @@
     "The pyarrow library is not installed, please install "
     "pyarrow to use the to_arrow() function."
 )
-_NO_TQDM_ERROR = (
-    "A progress bar was requested, but there was an error loading the tqdm "
-    "library. Please install tqdm to use the progress bar functionality."
-)
+
 _TABLE_HAS_NO_SCHEMA = 'Table has no schema:  call "client.get_table()"'
 
 
@@ -1371,32 +1364,6 @@ def total_rows(self):
         """int: The total number of rows in the table."""
         return self._total_rows
 
-    def _get_progress_bar(self, progress_bar_type):
-        """Construct a tqdm progress bar object, if tqdm is installed."""
-        if tqdm is None:
-            if progress_bar_type is not None:
-                warnings.warn(_NO_TQDM_ERROR, UserWarning, stacklevel=3)
-            return None
-
-        description = "Downloading"
-        unit = "rows"
-
-        try:
-            if progress_bar_type == "tqdm":
-                return tqdm.tqdm(desc=description, total=self.total_rows, unit=unit)
-            elif progress_bar_type == "tqdm_notebook":
-                return tqdm.tqdm_notebook(
-                    desc=description, total=self.total_rows, unit=unit
-                )
-            elif progress_bar_type == "tqdm_gui":
-                return tqdm.tqdm_gui(desc=description, total=self.total_rows, unit=unit)
-        except (KeyError, TypeError):
-            # Protect ourselves from any tqdm errors. In case of
-            # unexpected tqdm behavior, just fall back to showing
-            # no progress bar.
-            warnings.warn(_NO_TQDM_ERROR, UserWarning, stacklevel=3)
-        return None
-
     def _to_page_iterable(
         self, bqstorage_download, tabledata_list_download, bqstorage_client=None
     ):
@@ -1508,7 +1475,9 @@ def to_arrow(
             owns_bqstorage_client = bqstorage_client is not None
 
         try:
-            progress_bar = self._get_progress_bar(progress_bar_type)
+            progress_bar = _get_progress_bar(
+                progress_bar_type, "Downloading", self.total_rows, "rows"
+            )
 
             record_batches = []
             for record_batch in self._to_arrow_iterable(