Skip to content

Commit 8d207c7

Browse files
plamutplamut
committed
Autofetch table schema on load if not provided
1 parent 86bb5cf commit 8d207c7

File tree

2 files changed

+135
-5
lines changed

2 files changed

+135
-5
lines changed

‎bigquery/google/cloud/bigquery/client.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1531,6 +1531,29 @@ def load_table_from_dataframe(
15311531
if location is None:
15321532
location = self.location
15331533

1534+
# If table schema is not provided, we try to fetch the existing table
1535+
# schema, and check if dataframe schema is compatible with it.
1536+
if not job_config.schema:
1537+
try:
1538+
table = self.get_table(destination)
1539+
except google.api_core.exceptions.NotFound:
1540+
table = None
1541+
else:
1542+
table_col_names = {field.name for field in table.schema}
1543+
dframe_col_names = set(dataframe.columns)
1544+
1545+
in_dframe_only = dframe_col_names - table_col_names
1546+
if in_dframe_only:
1547+
raise ValueError(
1548+
"Dataframe contains columns that are not present in "
1549+
"table: {}".format(in_dframe_only)
1550+
)
1551+
1552+
# schema fields not present in the dataframe are not needed
1553+
job_config.schema = [
1554+
field for field in table.schema if field.name in dframe_col_names
1555+
]
1556+
15341557
job_config.schema = _pandas_helpers.dataframe_to_bq_schema(
15351558
dataframe, job_config.schema
15361559
)

‎bigquery/tests/unit/test_client.py‎

Lines changed: 112 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5228,15 +5228,23 @@ def test_load_table_from_file_bad_mode(self):
52285228
def test_load_table_from_dataframe(self):
52295229
from google.cloud.bigquery.client import _DEFAULT_NUM_RETRIES
52305230
from google.cloud.bigquery import job
5231+
from google.cloud.bigquery.schema import SchemaField
52315232

52325233
client = self._make_client()
52335234
records = [{"id": 1, "age": 100}, {"id": 2, "age": 60}]
52345235
dataframe = pandas.DataFrame(records)
52355236

5237+
get_table_patch = mock.patch(
5238+
"google.cloud.bigquery.client.Client.get_table",
5239+
autospec=True,
5240+
return_value=mock.Mock(
5241+
schema=[SchemaField("id", "INTEGER"), SchemaField("age", "INTEGER")]
5242+
),
5243+
)
52365244
load_patch = mock.patch(
52375245
"google.cloud.bigquery.client.Client.load_table_from_file", autospec=True
52385246
)
5239-
with load_patch as load_table_from_file:
5247+
with load_patch as load_table_from_file, get_table_patch:
52405248
client.load_table_from_dataframe(dataframe, self.TABLE_REF)
52415249

52425250
load_table_from_file.assert_called_once_with(
@@ -5263,15 +5271,23 @@ def test_load_table_from_dataframe(self):
52635271
def test_load_table_from_dataframe_w_client_location(self):
52645272
from google.cloud.bigquery.client import _DEFAULT_NUM_RETRIES
52655273
from google.cloud.bigquery import job
5274+
from google.cloud.bigquery.schema import SchemaField
52665275

52675276
client = self._make_client(location=self.LOCATION)
52685277
records = [{"id": 1, "age": 100}, {"id": 2, "age": 60}]
52695278
dataframe = pandas.DataFrame(records)
52705279

5280+
get_table_patch = mock.patch(
5281+
"google.cloud.bigquery.client.Client.get_table",
5282+
autospec=True,
5283+
return_value=mock.Mock(
5284+
schema=[SchemaField("id", "INTEGER"), SchemaField("age", "INTEGER")]
5285+
),
5286+
)
52715287
load_patch = mock.patch(
52725288
"google.cloud.bigquery.client.Client.load_table_from_file", autospec=True
52735289
)
5274-
with load_patch as load_table_from_file:
5290+
with load_patch as load_table_from_file, get_table_patch:
52755291
client.load_table_from_dataframe(dataframe, self.TABLE_REF)
52765292

52775293
load_table_from_file.assert_called_once_with(
@@ -5298,16 +5314,24 @@ def test_load_table_from_dataframe_w_client_location(self):
52985314
def test_load_table_from_dataframe_w_custom_job_config(self):
52995315
from google.cloud.bigquery.client import _DEFAULT_NUM_RETRIES
53005316
from google.cloud.bigquery import job
5317+
from google.cloud.bigquery.schema import SchemaField
53015318

53025319
client = self._make_client()
53035320
records = [{"id": 1, "age": 100}, {"id": 2, "age": 60}]
53045321
dataframe = pandas.DataFrame(records)
53055322
job_config = job.LoadJobConfig()
53065323

5324+
get_table_patch = mock.patch(
5325+
"google.cloud.bigquery.client.Client.get_table",
5326+
autospec=True,
5327+
return_value=mock.Mock(
5328+
schema=[SchemaField("id", "INTEGER"), SchemaField("age", "INTEGER")]
5329+
),
5330+
)
53075331
load_patch = mock.patch(
53085332
"google.cloud.bigquery.client.Client.load_table_from_file", autospec=True
53095333
)
5310-
with load_patch as load_table_from_file:
5334+
with load_patch as load_table_from_file, get_table_patch:
53115335
client.load_table_from_dataframe(
53125336
dataframe, self.TABLE_REF, job_config=job_config, location=self.LOCATION
53135337
)
@@ -5370,7 +5394,20 @@ def test_load_table_from_dataframe_w_automatic_schema(self):
53705394
"google.cloud.bigquery.client.Client.load_table_from_file", autospec=True
53715395
)
53725396

5373-
with load_patch as load_table_from_file:
5397+
get_table_patch = mock.patch(
5398+
"google.cloud.bigquery.client.Client.get_table",
5399+
autospec=True,
5400+
return_value=mock.Mock(
5401+
schema=[
5402+
SchemaField("int_col", "INTEGER"),
5403+
SchemaField("float_col", "FLOAT"),
5404+
SchemaField("bool_col", "BOOLEAN"),
5405+
SchemaField("dt_col", "DATETIME"),
5406+
SchemaField("ts_col", "TIMESTAMP"),
5407+
]
5408+
),
5409+
)
5410+
with load_patch as load_table_from_file, get_table_patch:
53745411
client.load_table_from_dataframe(
53755412
dataframe, self.TABLE_REF, location=self.LOCATION
53765413
)
@@ -5398,6 +5435,67 @@ def test_load_table_from_dataframe_w_automatic_schema(self):
53985435
SchemaField("ts_col", "TIMESTAMP"),
53995436
)
54005437

5438+
@unittest.skipIf(pandas is None, "Requires `pandas`")
5439+
@unittest.skipIf(pyarrow is None, "Requires `pyarrow`")
5440+
def test_load_table_from_dataframe_unknown_df_columns(self):
5441+
from google.cloud.bigquery.schema import SchemaField
5442+
5443+
client = self._make_client()
5444+
records = [{"id": 1, "typo_age": 100}, {"id": 2, "typo_age": 60}]
5445+
dataframe = pandas.DataFrame(records)
5446+
5447+
get_table_patch = mock.patch(
5448+
"google.cloud.bigquery.client.Client.get_table",
5449+
autospec=True,
5450+
return_value=mock.Mock(
5451+
schema=[SchemaField("id", "INTEGER"), SchemaField("age", "INTEGER")]
5452+
),
5453+
)
5454+
load_patch = mock.patch(
5455+
"google.cloud.bigquery.client.Client.load_table_from_file", autospec=True
5456+
)
5457+
with pytest.raises(ValueError) as exc_info, load_patch, get_table_patch:
5458+
client.load_table_from_dataframe(dataframe, self.TABLE_REF)
5459+
5460+
err_msg = str(exc_info.value)
5461+
assert "Dataframe contains columns that are not present in table" in err_msg
5462+
assert "typo_age" in err_msg
5463+
assert "id" not in err_msg
5464+
5465+
@unittest.skipIf(pandas is None, "Requires `pandas`")
5466+
@unittest.skipIf(pyarrow is None, "Requires `pyarrow`")
5467+
def test_load_table_from_dataframe_unknown_table(self):
5468+
from google.cloud.bigquery.client import _DEFAULT_NUM_RETRIES
5469+
5470+
client = self._make_client()
5471+
records = [{"id": 1, "age": 100}, {"id": 2, "age": 60}]
5472+
dataframe = pandas.DataFrame(records)
5473+
5474+
get_table_patch = mock.patch(
5475+
"google.cloud.bigquery.client.Client.get_table",
5476+
autospec=True,
5477+
side_effect=google.api_core.exceptions.NotFound("Table not found"),
5478+
)
5479+
load_patch = mock.patch(
5480+
"google.cloud.bigquery.client.Client.load_table_from_file", autospec=True
5481+
)
5482+
with load_patch as load_table_from_file, get_table_patch:
5483+
# there should be no error
5484+
client.load_table_from_dataframe(dataframe, self.TABLE_REF)
5485+
5486+
load_table_from_file.assert_called_once_with(
5487+
client,
5488+
mock.ANY,
5489+
self.TABLE_REF,
5490+
num_retries=_DEFAULT_NUM_RETRIES,
5491+
rewind=True,
5492+
job_id=mock.ANY,
5493+
job_id_prefix=None,
5494+
location=None,
5495+
project=None,
5496+
job_config=mock.ANY,
5497+
)
5498+
54015499
@unittest.skipIf(pandas is None, "Requires `pandas`")
54025500
@unittest.skipIf(pyarrow is None, "Requires `pyarrow`")
54035501
def test_load_table_from_dataframe_struct_fields_error(self):
@@ -5686,10 +5784,19 @@ def test_load_table_from_dataframe_w_schema_arrow_custom_compression(self):
56865784
@unittest.skipIf(pandas is None, "Requires `pandas`")
56875785
@unittest.skipIf(pyarrow is None, "Requires `pyarrow`")
56885786
def test_load_table_from_dataframe_wo_pyarrow_custom_compression(self):
5787+
from google.cloud.bigquery.schema import SchemaField
5788+
56895789
client = self._make_client()
56905790
records = [{"id": 1, "age": 100}, {"id": 2, "age": 60}]
56915791
dataframe = pandas.DataFrame(records)
56925792

5793+
get_table_patch = mock.patch(
5794+
"google.cloud.bigquery.client.Client.get_table",
5795+
autospec=True,
5796+
return_value=mock.Mock(
5797+
schema=[SchemaField("id", "INTEGER"), SchemaField("age", "INTEGER")]
5798+
),
5799+
)
56935800
load_patch = mock.patch(
56945801
"google.cloud.bigquery.client.Client.load_table_from_file", autospec=True
56955802
)
@@ -5698,7 +5805,7 @@ def test_load_table_from_dataframe_wo_pyarrow_custom_compression(self):
56985805
dataframe, "to_parquet", wraps=dataframe.to_parquet
56995806
)
57005807

5701-
with load_patch, pyarrow_patch, to_parquet_patch as to_parquet_spy:
5808+
with load_patch, get_table_patch, pyarrow_patch, to_parquet_patch as to_parquet_spy:
57025809
client.load_table_from_dataframe(
57035810
dataframe,
57045811
self.TABLE_REF,

0 commit comments

Comments
 (0)