This is an automated email from the ASF dual-hosted git repository.
claudevdm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new d71149317eb [Python] Reset BigQuery retries for each batch (#39995)
d71149317eb is described below
commit d71149317ebf46376968e16367d8c4f8f6e5b0ea
Author: Bruno Volpato <[email protected]>
AuthorDate: Mon Oct 12 08:19:35 2026 -0400
[Python] Reset BigQuery retries for each batch (#39995)
* [Python] Reset BigQuery retries for each batch
* [Python] Link BigQuery retry fix in changelog
* [Python] Move BigQuery retry fix to unreleased changelog
---
CHANGES.md | 1 +
sdks/python/apache_beam/io/gcp/bigquery.py | 14 ++--
sdks/python/apache_beam/io/gcp/bigquery_test.py | 104 ++++++++++++++++++++++++
3 files changed, 112 insertions(+), 7 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index cb5eced42f7..6b01d163823 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -89,6 +89,7 @@
## Bugfixes
+* (Python) BigQuery streaming inserts now reset the retry budget and backoff
for each batch instead of sharing them across a bundle
([#39995](https://github.com/apache/beam/pull/39995)).
* (Go) Fixed a data race on the Prism runner's artifact cache map in
JobServices ([#32656](https://github.com/apache/beam/issues/32656)).
* (Go) Fixed the harness leaking Data/State gRPC streams after the worker
stops, and a deadlock when Send returns EOF
([#40260](https://github.com/apache/beam/issues/40260)).
* (Go) Fixed a deadlock recreating a data channel while holding the channel
lock ([#40414](https://github.com/apache/beam/issues/40414)).
diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py
b/sdks/python/apache_beam/io/gcp/bigquery.py
index f6dd781695c..2a52b31edf7 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery.py
@@ -1640,12 +1640,6 @@ class BigQueryWriteFn(DoFn):
minimum_logging_frequency_msec
) = self.streaming_api_logging_frequency_sec * 1000
- self._backoff_calculator = iter(
- retry.FuzzedExponentialIntervals(
- initial_delay_secs=0.2,
- num_retries=self._max_retries,
- max_delay_secs=1500))
-
def _create_table_if_needed(self, table_reference, schema=None):
str_table_reference = '%s:%s.%s' % (
table_reference.projectId,
@@ -1824,6 +1818,12 @@ class BigQueryWriteFn(DoFn):
else:
insert_ids = [r[1] for r in rows_and_insert_ids]
+ backoff_calculator = iter(
+ retry.FuzzedExponentialIntervals(
+ initial_delay_secs=0.2,
+ num_retries=self._max_retries,
+ max_delay_secs=1500))
+
while True:
start = time.time()
passed, errors = self.bigquery_wrapper.insert_rows(
@@ -1840,7 +1840,7 @@ class BigQueryWriteFn(DoFn):
rows[entry['index']], entry["errors"], window_values[entry['index']])
for entry in errors]
failed_insert_ids = [insert_ids[entry['index']] for entry in errors]
- retry_backoff = next(self._backoff_calculator, None)
+ retry_backoff = next(backoff_calculator, None)
# If retry_backoff is None, then we will not retry and must log.
should_retry = any(
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_test.py
b/sdks/python/apache_beam/io/gcp/bigquery_test.py
index d7d4cab09d9..22b617fa874 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_test.py
@@ -2335,6 +2335,110 @@ def test_with_batched_input_splits_large_batch(self):
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class BigQueryStreamingInsertTransformTests(unittest.TestCase):
+ def _retry_test_writer(self, responses, **kwargs):
+ client = mock.Mock()
+ client.insert_rows_json.side_effect = responses
+ fn = beam_bq.BigQueryWriteFn(
+ batch_size=1,
+ create_disposition=beam_bq.BigQueryDisposition.CREATE_NEVER,
+ write_disposition=beam_bq.BigQueryDisposition.WRITE_APPEND,
+ with_batched_input=True,
+ test_client=client,
+ **kwargs)
+ fn.start_bundle()
+ return fn, client
+
+ def _process_retry_batch(
+ self, fn, month, destination='project-id:dataset_id.table_id'):
+ return fn.process(
+ (destination, [({
+ 'month': month
+ }, 'insertid%d' % month)]),
+ window_value=beam.transforms.window.GlobalWindows.windowed_value(None))
+
+ @parameterized.expand([
+ ('same_destination', 'project-id:dataset_id.table_id'),
+ ('different_destination', 'project-id:dataset_id.other_table'),
+ ])
+ @mock.patch('apache_beam.io.gcp.bigquery.time.sleep')
+ def test_successful_batches_do_not_consume_retries(
+ self, unused_name, destination, sleep):
+ errors = [{'index': 0, 'errors': [{'reason': 'backendError'}]}]
+ fn, client = self._retry_test_writer([[], errors, []], max_retries=1)
+
+ self.assertEqual([], self._process_retry_batch(fn, 1))
+ self.assertEqual([], self._process_retry_batch(fn, 2, destination))
+
+ self.assertEqual([[{
+ 'month': 1
+ }], [{
+ 'month': 2
+ }], [{
+ 'month': 2
+ }]],
+ [
+ call.kwargs['json_rows']
+ for call in client.insert_rows_json.call_args_list
+ ])
+ self.assertEqual([['insertid1'], ['insertid2'], ['insertid2']],
+ [
+ call.kwargs['row_ids']
+ for call in client.insert_rows_json.call_args_list
+ ])
+ sleep.assert_called_once()
+
+ @mock.patch('apache_beam.io.gcp.bigquery.time.sleep')
+ def test_exhausted_batch_does_not_consume_next_batch_retries(self, sleep):
+ errors = [{'index': 0, 'errors': [{'reason': 'backendError'}]}]
+ fn, client = self._retry_test_writer(
+ [errors, errors, errors, []], max_retries=1)
+
+ failed = self._process_retry_batch(fn, 1)
+ self.assertEqual([('project-id:dataset_id.table_id', {
+ 'month': 1
+ })],
+ [
+ output.value.value for output in failed
+ if output.tag == beam_bq.BigQueryWriteFn.FAILED_ROWS
+ ])
+ self.assertEqual([], self._process_retry_batch(fn, 2))
+ self.assertEqual(4, client.insert_rows_json.call_count)
+ self.assertEqual(2, sleep.call_count)
+
+ @mock.patch('apache_beam.io.gcp.bigquery.time.sleep')
+ def test_successful_batches_do_not_increase_retry_backoff(self, sleep):
+ errors = [{'index': 0, 'errors': [{'reason': 'backendError'}]}]
+ fn, client = self._retry_test_writer([[]] * 20 + [errors, errors, []])
+
+ for month in range(21):
+ self.assertEqual([], self._process_retry_batch(fn, month))
+
+ self.assertEqual(23, client.insert_rows_json.call_count)
+ self.assertEqual(2, sleep.call_count)
+ first_delay = sleep.call_args_list[0].args[0]
+ second_delay = sleep.call_args_list[1].args[0]
+ self.assertGreaterEqual(first_delay, 0.1)
+ self.assertLessEqual(first_delay, 0.2)
+ self.assertGreaterEqual(second_delay, 0.2)
+ self.assertLessEqual(second_delay, 0.4)
+
+ @mock.patch('apache_beam.io.gcp.bigquery.time.sleep')
+ def test_zero_retries_attempts_each_batch_once(self, sleep):
+ errors = [{'index': 0, 'errors': [{'reason': 'backendError'}]}]
+ fn, client = self._retry_test_writer([errors, []], max_retries=0)
+
+ failed = self._process_retry_batch(fn, 1)
+ self.assertEqual([('project-id:dataset_id.table_id', {
+ 'month': 1
+ })],
+ [
+ output.value.value for output in failed
+ if output.tag == beam_bq.BigQueryWriteFn.FAILED_ROWS
+ ])
+ self.assertEqual([], self._process_retry_batch(fn, 2))
+ self.assertEqual(2, client.insert_rows_json.call_count)
+ sleep.assert_not_called()
+
def test_dofn_client_process_performs_batching(self):
client = mock.Mock()
client.tables.Get.return_value = bigquery.Table(
