Skip to content

Actually use the progress_bar_type argument in QueryJob.to_dataframe() #7616

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tseaver merged 3 commits into from
Mar 30, 2019
Merged

Actually use the progress_bar_type argument in QueryJob.to_dataframe() #7616

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
89fccc1
Actually use the progress_bar_type argument in QueryJob.to_dataframe()
tswast Mar 29, 2019
0b73e32
Test that a progress bar is constructed from QueryJob.to_dataframe
tswast Mar 29, 2019
3b1d0d4
Blacken
tswast Mar 29, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion bigquery/google/cloud/bigquery/job.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
8 changes: 4 additions & 4 deletions bigquery/google/cloud/bigquery/table.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -1515,14 +1513,16 @@ 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:
bqstorage_client (Any):
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:
Expand Down
35 changes: 35 additions & 0 deletions bigquery/tests/unit/test_job.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down Expand Up @@ -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

Expand Down