Skip to content

Commit 501e993

Browse files
committed
Refactor: Use generator for request creation and block on send
- Changed `generate_write_requests` to be a generator, yielding requests instead of returning a list. - Made `stream.send()` calls blocking by calling `future.result()` immediately, ensuring requests are sent sequentially.
1 parent 610f38e commit 501e993

File tree

1 file changed

+4
-10
lines changed

1 file changed

+4
-10
lines changed

packages/google-cloud-bigquery-storage/samples/pyarrow/append_rows_with_arrow.py

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -166,7 +166,6 @@ def generate_write_requests(pyarrow_table):
166166
# Maximum size for a single AppendRowsRequest is 10 MB.
167167
# To be safe, we'll aim for a soft limit of 7 MB.
168168
max_request_bytes = 7 * 1024 * 1024 # 7 MB
169-
requests = []
170169

171170
def _create_request(batches):
172171
"""Helper to create an AppendRowsRequest from a list of batches."""
@@ -214,8 +213,7 @@ def _create_request(batches):
214213
)
215214
# otherwise, generate the request, reset current_size and current_batches
216215
else:
217-
request = _create_request(current_batches)
218-
requests.append(request)
216+
yield _create_request(current_batches)
219217

220218
current_batches = []
221219
current_size = 0
@@ -228,10 +226,7 @@ def _create_request(batches):
228226

229227
# Flush remaining batches
230228
if current_batches:
231-
request = _create_request(current_batches)
232-
requests.append(request)
233-
234-
return requests
229+
yield _create_request(current_batches)
235230

236231

237232
def verify_result(client, table, futures):
@@ -272,8 +267,7 @@ def main(project_id, dataset):
272267
for request in requests:
273268
future = stream.send(request)
274269
futures.append(future)
275-
# future.result() # Optional, will block until writing is complete.
276-
for future in futures:
277-
future.result()
270+
future.result() # Optional, will block until writing is complete.
271+
278272
# Verify results.
279273
verify_result(bq_client, bq_table, futures)

0 commit comments

Comments
 (0)