diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py index bc87f109a484..94a2290cc29e 100644 --- a/bigquery/google/cloud/bigquery/job.py +++ b/bigquery/google/cloud/bigquery/job.py @@ -2857,7 +2857,9 @@ def to_dataframe(self, bqstorage_client=None, dtypes=None, progress_bar_type=Non ValueError: If the `pandas` library cannot be imported. """ return self.result().to_dataframe( - bqstorage_client=bqstorage_client, dtypes=dtypes + bqstorage_client=bqstorage_client, + dtypes=dtypes, + progress_bar_type=progress_bar_type, ) def __iter__(self): diff --git a/bigquery/google/cloud/bigquery/table.py b/bigquery/google/cloud/bigquery/table.py index dcb25d8bb3c6..5b0ad45af4be 100644 --- a/bigquery/google/cloud/bigquery/table.py +++ b/bigquery/google/cloud/bigquery/table.py @@ -1422,9 +1422,7 @@ def _get_progress_bar(self, progress_bar_type): 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 - ) + 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 @@ -1515,7 +1513,7 @@ class _EmptyRowIterator(object): pages = () total_rows = 0 - def to_dataframe(self, bqstorage_client=None, dtypes=None): + def to_dataframe(self, bqstorage_client=None, dtypes=None, progress_bar_type=None): """Create an empty dataframe. Args: @@ -1523,6 +1521,8 @@ def to_dataframe(self, bqstorage_client=None, dtypes=None): Ignored. Added for compatibility with RowIterator. dtypes (Any): Ignored. Added for compatibility with RowIterator. + progress_bar_type (Any): + Ignored. Added for compatibility with RowIterator. Returns: pandas.DataFrame: diff --git a/bigquery/tests/unit/test_job.py b/bigquery/tests/unit/test_job.py index 833081ce066d..a42d9ffc311c 100644 --- a/bigquery/tests/unit/test_job.py +++ b/bigquery/tests/unit/test_job.py @@ -27,6 +27,10 @@ from google.cloud import bigquery_storage_v1beta1 except (ImportError, AttributeError): # pragma: NO COVER bigquery_storage_v1beta1 = None +try: + from tqdm import tqdm +except (ImportError, AttributeError): # pragma: NO COVER + tqdm = None def _make_credentials(): @@ -4699,6 +4703,37 @@ def test_to_dataframe_column_dtypes(self): self.assertEqual(df.complete.dtype.name, "bool") self.assertEqual(df.date.dtype.name, "object") + @unittest.skipIf(pandas is None, "Requires `pandas`") + @unittest.skipIf(tqdm is None, "Requires `tqdm`") + @mock.patch("tqdm.tqdm") + def test_to_dataframe_with_progress_bar(self, tqdm_mock): + begun_resource = self._make_resource() + query_resource = { + "jobComplete": True, + "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID}, + "totalRows": "4", + "schema": { + "fields": [{"name": "name", "type": "STRING", "mode": "NULLABLE"}] + }, + } + done_resource = copy.deepcopy(begun_resource) + done_resource["status"] = {"state": "DONE"} + connection = _make_connection( + begun_resource, + query_resource, + done_resource, + query_resource, + query_resource, + ) + client = _make_client(project=self.PROJECT, connection=connection) + job = self._make_one(self.JOB_ID, self.QUERY, client) + + job.to_dataframe(progress_bar_type=None) + tqdm_mock.assert_not_called() + + job.to_dataframe(progress_bar_type="tqdm") + tqdm_mock.assert_called() + def test_iter(self): import types