Skip to content

Commit 3be78b7

Browse files
perf: add size parameter for load table from dataframe and json methods (#280)
* feat: add size parameter for load from dataframe and json * pref: calculate length of encoded string
1 parent c69cd50 commit 3be78b7

File tree

2 files changed

+20
-2
lines changed

2 files changed

+20
-2
lines changed

google/cloud/bigquery/client.py

+5-2
Original file line numberDiff line numberDiff line change
@@ -2237,11 +2237,13 @@ def load_table_from_dataframe(
22372237
dataframe.to_parquet(tmppath, compression=parquet_compression)
22382238

22392239
with open(tmppath, "rb") as parquet_file:
2240+
file_size = os.path.getsize(tmppath)
22402241
return self.load_table_from_file(
22412242
parquet_file,
22422243
destination,
22432244
num_retries=num_retries,
22442245
rewind=True,
2246+
size=file_size,
22452247
job_id=job_id,
22462248
job_id_prefix=job_id_prefix,
22472249
location=location,
@@ -2343,11 +2345,12 @@ def load_table_from_json(
23432345
destination = _table_arg_to_table_ref(destination, default_project=self.project)
23442346

23452347
data_str = u"\n".join(json.dumps(item) for item in json_rows)
2346-
data_file = io.BytesIO(data_str.encode())
2347-
2348+
encoded_str = data_str.encode()
2349+
data_file = io.BytesIO(encoded_str)
23482350
return self.load_table_from_file(
23492351
data_file,
23502352
destination,
2353+
size=len(encoded_str),
23512354
num_retries=num_retries,
23522355
job_id=job_id,
23532356
job_id_prefix=job_id_prefix,

tests/unit/test_client.py

+15
Original file line numberDiff line numberDiff line change
@@ -7482,6 +7482,7 @@ def test_load_table_from_dataframe(self):
74827482
self.TABLE_REF,
74837483
num_retries=_DEFAULT_NUM_RETRIES,
74847484
rewind=True,
7485+
size=mock.ANY,
74857486
job_id=mock.ANY,
74867487
job_id_prefix=None,
74877488
location=None,
@@ -7525,6 +7526,7 @@ def test_load_table_from_dataframe_w_client_location(self):
75257526
self.TABLE_REF,
75267527
num_retries=_DEFAULT_NUM_RETRIES,
75277528
rewind=True,
7529+
size=mock.ANY,
75287530
job_id=mock.ANY,
75297531
job_id_prefix=None,
75307532
location=self.LOCATION,
@@ -7577,6 +7579,7 @@ def test_load_table_from_dataframe_w_custom_job_config_wihtout_source_format(sel
75777579
self.TABLE_REF,
75787580
num_retries=_DEFAULT_NUM_RETRIES,
75797581
rewind=True,
7582+
size=mock.ANY,
75807583
job_id=mock.ANY,
75817584
job_id_prefix=None,
75827585
location=self.LOCATION,
@@ -7631,6 +7634,7 @@ def test_load_table_from_dataframe_w_custom_job_config_w_source_format(self):
76317634
self.TABLE_REF,
76327635
num_retries=_DEFAULT_NUM_RETRIES,
76337636
rewind=True,
7637+
size=mock.ANY,
76347638
job_id=mock.ANY,
76357639
job_id_prefix=None,
76367640
location=self.LOCATION,
@@ -7723,6 +7727,7 @@ def test_load_table_from_dataframe_w_automatic_schema(self):
77237727
self.TABLE_REF,
77247728
num_retries=_DEFAULT_NUM_RETRIES,
77257729
rewind=True,
7730+
size=mock.ANY,
77267731
job_id=mock.ANY,
77277732
job_id_prefix=None,
77287733
location=self.LOCATION,
@@ -7782,6 +7787,7 @@ def test_load_table_from_dataframe_w_index_and_auto_schema(self):
77827787
self.TABLE_REF,
77837788
num_retries=_DEFAULT_NUM_RETRIES,
77847789
rewind=True,
7790+
size=mock.ANY,
77857791
job_id=mock.ANY,
77867792
job_id_prefix=None,
77877793
location=self.LOCATION,
@@ -7827,6 +7833,7 @@ def test_load_table_from_dataframe_unknown_table(self):
78277833
self.TABLE_REF,
78287834
num_retries=_DEFAULT_NUM_RETRIES,
78297835
rewind=True,
7836+
size=mock.ANY,
78307837
job_id=mock.ANY,
78317838
job_id_prefix=None,
78327839
location=None,
@@ -7867,6 +7874,7 @@ def test_load_table_from_dataframe_w_nullable_int64_datatype(self):
78677874
self.TABLE_REF,
78687875
num_retries=_DEFAULT_NUM_RETRIES,
78697876
rewind=True,
7877+
size=mock.ANY,
78707878
job_id=mock.ANY,
78717879
job_id_prefix=None,
78727880
location=self.LOCATION,
@@ -7913,6 +7921,7 @@ def test_load_table_from_dataframe_w_nullable_int64_datatype_automatic_schema(se
79137921
self.TABLE_REF,
79147922
num_retries=_DEFAULT_NUM_RETRIES,
79157923
rewind=True,
7924+
size=mock.ANY,
79167925
job_id=mock.ANY,
79177926
job_id_prefix=None,
79187927
location=self.LOCATION,
@@ -7973,6 +7982,7 @@ def test_load_table_from_dataframe_struct_fields(self):
79737982
self.TABLE_REF,
79747983
num_retries=_DEFAULT_NUM_RETRIES,
79757984
rewind=True,
7985+
size=mock.ANY,
79767986
job_id=mock.ANY,
79777987
job_id_prefix=None,
79787988
location=self.LOCATION,
@@ -8046,6 +8056,7 @@ def test_load_table_from_dataframe_w_partial_schema(self):
80468056
self.TABLE_REF,
80478057
num_retries=_DEFAULT_NUM_RETRIES,
80488058
rewind=True,
8059+
size=mock.ANY,
80498060
job_id=mock.ANY,
80508061
job_id_prefix=None,
80518062
location=self.LOCATION,
@@ -8139,6 +8150,7 @@ def test_load_table_from_dataframe_w_partial_schema_missing_types(self):
81398150
self.TABLE_REF,
81408151
num_retries=_DEFAULT_NUM_RETRIES,
81418152
rewind=True,
8153+
size=mock.ANY,
81428154
job_id=mock.ANY,
81438155
job_id_prefix=None,
81448156
location=self.LOCATION,
@@ -8251,6 +8263,7 @@ def test_load_table_from_dataframe_w_nulls(self):
82518263
self.TABLE_REF,
82528264
num_retries=_DEFAULT_NUM_RETRIES,
82538265
rewind=True,
8266+
size=mock.ANY,
82548267
job_id=mock.ANY,
82558268
job_id_prefix=None,
82568269
location=self.LOCATION,
@@ -8302,6 +8315,7 @@ def test_load_table_from_json_basic_use(self):
83028315
client,
83038316
mock.ANY,
83048317
self.TABLE_REF,
8318+
size=mock.ANY,
83058319
num_retries=_DEFAULT_NUM_RETRIES,
83068320
job_id=mock.ANY,
83078321
job_id_prefix=None,
@@ -8353,6 +8367,7 @@ def test_load_table_from_json_non_default_args(self):
83538367
client,
83548368
mock.ANY,
83558369
self.TABLE_REF,
8370+
size=mock.ANY,
83568371
num_retries=_DEFAULT_NUM_RETRIES,
83578372
job_id=mock.ANY,
83588373
job_id_prefix=None,

0 commit comments

Comments
 (0)