Skip to content

Commit d038ace

Browse files
aktechtswast
authored andcommitted
Add progress for to_gbq function (#166)
* Add progress for to_gbq function using tqdm * Add tqdm addition to changelog * Add issue number to changelog for progress bar
1 parent 7711bb0 commit d038ace

File tree

4 files changed

+24
-6
lines changed

4 files changed

+24
-6
lines changed

docs/source/changelog.rst

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,9 @@ Changelog
77
- Project ID parameter is optional in ``read_gbq`` and ``to_gbq`` when it can
88
inferred from the environment. Note: you must still pass in a project ID when
99
using user-based authentication. (:issue:`103`)
10+
- Progress bar added for ``to_gbq``, through an optional library `tqdm` as
11+
dependency. (:issue:`162`)
12+
1013

1114
Internal changes
1215
~~~~~~~~~~~~~~~~

pandas_gbq/gbq.py

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,11 @@
1616
BIGQUERY_INSTALLED_VERSION = None
1717
SHOW_VERBOSE_DEPRECATION = False
1818

19+
try:
20+
import tqdm # noqa
21+
except ImportError:
22+
tqdm = None
23+
1924

2025
def _check_google_client_version():
2126
global BIGQUERY_INSTALLED_VERSION, SHOW_VERBOSE_DEPRECATION
@@ -563,16 +568,19 @@ def run_query(self, query, **kwargs):
563568

564569
def load_data(
565570
self, dataframe, dataset_id, table_id, chunksize=None,
566-
schema=None):
571+
schema=None, progress_bar=True):
567572
from pandas_gbq import load
568573

569574
total_rows = len(dataframe)
570575
logger.info("\n\n")
571576

572577
try:
573-
for remaining_rows in load.load_chunks(
574-
self.client, dataframe, dataset_id, table_id,
575-
chunksize=chunksize, schema=schema):
578+
chunks = load.load_chunks(self.client, dataframe, dataset_id,
579+
table_id, chunksize=chunksize,
580+
schema=schema)
581+
if progress_bar and tqdm:
582+
chunks = tqdm.tqdm(chunks)
583+
for remaining_rows in chunks:
576584
logger.info("\rLoad is {0}% Complete".format(
577585
((total_rows - remaining_rows) * 100) / total_rows))
578586
except self.http_error as ex:
@@ -870,7 +878,7 @@ def read_gbq(query, project_id=None, index_col=None, col_order=None,
870878

871879
def to_gbq(dataframe, destination_table, project_id=None, chunksize=None,
872880
verbose=None, reauth=False, if_exists='fail', private_key=None,
873-
auth_local_webserver=False, table_schema=None):
881+
auth_local_webserver=False, table_schema=None, progress_bar=True):
874882
"""Write a DataFrame to a Google BigQuery table.
875883
876884
The main method a user calls to export pandas DataFrame contents to
@@ -935,6 +943,8 @@ def to_gbq(dataframe, destination_table, project_id=None, chunksize=None,
935943
names of a field.
936944
.. versionadded:: 0.3.1
937945
verbose : None, deprecated
946+
progress_bar : boolean, True by default. It uses the library `tqdm` to show
947+
the progress bar for the upload, chunk by chunk.
938948
"""
939949

940950
_test_google_api_imports()
@@ -987,7 +997,7 @@ def to_gbq(dataframe, destination_table, project_id=None, chunksize=None,
987997

988998
connector.load_data(
989999
dataframe, dataset_id, table_id, chunksize=chunksize,
990-
schema=table_schema)
1000+
schema=table_schema, progress_bar=progress_bar)
9911001

9921002

9931003
def generate_bq_schema(df, default_type='STRING'):

requirements.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,3 +2,4 @@ pandas
22
google-auth
33
google-auth-oauthlib
44
google-cloud-bigquery
5+
tqdm

setup.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,9 @@ def readme():
2424
'google-cloud-bigquery>=0.29.0',
2525
]
2626

27+
extras = {
28+
'tqdm': 'tqdm>=4.23.0',
29+
}
2730

2831
setup(
2932
name=NAME,
@@ -50,6 +53,7 @@ def readme():
5053
],
5154
keywords='data',
5255
install_requires=INSTALL_REQUIRES,
56+
extras_require=extras,
5357
packages=find_packages(exclude=['contrib', 'docs', 'tests*']),
5458
test_suite='tests',
5559
)

0 commit comments

Comments
 (0)