Skip to content

Commit 6792076

Browse files
mremesMatti Remes
authored and
Matti Remes
committed
ENH: Add table_schema parameter for user-defined BigQuery schema (#46)
1 parent 336fd78 commit 6792076

File tree

3 files changed

+55
-4
lines changed

3 files changed

+55
-4
lines changed

docs/source/changelog.rst

+1-2
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,11 @@
11
Changelog
22
=========
33

4-
54
0.3.1 / [TBD]
65
------------------
76

87
- Fix an issue where Unicode couldn't be uploaded in Python 2 (:issue:`93`)
9-
8+
- Add support for a passed schema in :func:``to_gbq`` instead inferring the schema from the passed ``DataFrame`` with ``DataFrame.dtypes`` (:issue:`46`)
109

1110
0.3.0 / 2018-01-03
1211
------------------

pandas_gbq/gbq.py

+12-2
Original file line numberDiff line numberDiff line change
@@ -891,7 +891,7 @@ def read_gbq(query, project_id=None, index_col=None, col_order=None,
891891

892892
def to_gbq(dataframe, destination_table, project_id, chunksize=10000,
893893
verbose=True, reauth=False, if_exists='fail', private_key=None,
894-
auth_local_webserver=False):
894+
auth_local_webserver=False, table_schema=None):
895895
"""Write a DataFrame to a Google BigQuery table.
896896
897897
The main method a user calls to export pandas DataFrame contents to
@@ -949,6 +949,13 @@ def to_gbq(dataframe, destination_table, project_id, chunksize=10000,
949949
.. [console flow]
950950
http://google-auth-oauthlib.readthedocs.io/en/latest/reference/google_auth_oauthlib.flow.html#google_auth_oauthlib.flow.InstalledAppFlow.run_console
951951
.. versionadded:: 0.2.0
952+
table_schema : list of dicts
953+
List of BigQuery table fields to which according DataFrame columns
954+
conform to, e.g. `[{'name': 'col1', 'type': 'STRING'},...]`. If
955+
schema is not provided, it will be generated according to dtypes
956+
of DataFrame columns. See BigQuery API documentation on available
957+
names of a field.
958+
.. versionadded:: 0.3.0
952959
"""
953960

954961
_test_google_api_imports()
@@ -968,7 +975,10 @@ def to_gbq(dataframe, destination_table, project_id, chunksize=10000,
968975
table = _Table(project_id, dataset_id, reauth=reauth,
969976
private_key=private_key)
970977

971-
table_schema = _generate_bq_schema(dataframe)
978+
if not table_schema:
979+
table_schema = _generate_bq_schema(dataframe)
980+
else:
981+
table_schema = dict(fields=table_schema)
972982

973983
# If table exists, check if_exists parameter
974984
if table.exists(table_id):

pandas_gbq/tests/test_gbq.py

+42
Original file line numberDiff line numberDiff line change
@@ -1422,6 +1422,48 @@ def test_schema_is_subset_fails_if_not_subset(self):
14221422
assert self.sut.schema_is_subset(
14231423
dataset, table_name, tested_schema) is False
14241424

1425+
def test_upload_data_with_valid_user_schema(self):
1426+
# Issue #46; tests test scenarios with user-provided
1427+
# schemas
1428+
df = tm.makeMixedDataFrame()
1429+
test_id = "15"
1430+
test_schema = [{'name': 'A', 'type': 'FLOAT'},
1431+
{'name': 'B', 'type': 'FLOAT'},
1432+
{'name': 'C', 'type': 'STRING'},
1433+
{'name': 'D', 'type': 'TIMESTAMP'}]
1434+
destination_table = self.destination_table + test_id
1435+
gbq.to_gbq(df, destination_table, _get_project_id(),
1436+
private_key=_get_private_key_path(),
1437+
table_schema=test_schema)
1438+
dataset, table = destination_table.split('.')
1439+
assert self.table.verify_schema(dataset, table,
1440+
dict(fields=test_schema))
1441+
1442+
def test_upload_data_with_invalid_user_schema_raises_error(self):
1443+
df = tm.makeMixedDataFrame()
1444+
test_id = "16"
1445+
test_schema = [{'name': 'A', 'type': 'FLOAT'},
1446+
{'name': 'B', 'type': 'FLOAT'},
1447+
{'name': 'C', 'type': 'FLOAT'},
1448+
{'name': 'D', 'type': 'FLOAT'}]
1449+
destination_table = self.destination_table + test_id
1450+
with tm.assertRaises(gbq.StreamingInsertError):
1451+
gbq.to_gbq(df, destination_table, _get_project_id(),
1452+
private_key=_get_private_key_path(),
1453+
table_schema=test_schema)
1454+
1455+
def test_upload_data_with_missing_schema_fields_raises_error(self):
1456+
df = tm.makeMixedDataFrame()
1457+
test_id = "16"
1458+
test_schema = [{'name': 'A', 'type': 'FLOAT'},
1459+
{'name': 'B', 'type': 'FLOAT'},
1460+
{'name': 'C', 'type': 'FLOAT'}]
1461+
destination_table = self.destination_table + test_id
1462+
with tm.assertRaises(gbq.StreamingInsertError):
1463+
gbq.to_gbq(df, destination_table, _get_project_id(),
1464+
private_key=_get_private_key_path(),
1465+
table_schema=test_schema)
1466+
14251467
def test_list_dataset(self):
14261468
dataset_id = self.dataset_prefix + "1"
14271469
assert dataset_id in self.dataset.datasets()

0 commit comments

Comments
 (0)