From acaace20ca210f1cb799d2b4dbff590e291cfd4c Mon Sep 17 00:00:00 2001 From: Valentyn Tymofieiev Date: Wed, 4 Nov 2026 17:07:04 -0800 Subject: [PATCH] Fix flaky PipelineBasedStreamingInsertTest.test_failure_in_some_rows_does_not_duplicate The mocked insert_rows_json call failed rows by their position in the batch (index 0 and 2), assuming they are always 'value1' and 'value5'. However, _StreamToBigQuery assigns each row a random shard key in [0, 500] and applies ReshufflePerKey, so the order in which rows reach the client is not guaranteed. When 'value1' and 'value5' land on the same shard (~0.2% of runs), the rows arrive as [value1, value5, value3], and 'value3' is failed instead of 'value5', which breaks the RETRY_NEVER variant. Identify the rows to fail by value instead of by index. --- .../apache_beam/io/gcp/bigquery_test.py | 56 +++++++++++-------- 1 file changed, 33 insertions(+), 23 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_test.py b/sdks/python/apache_beam/io/gcp/bigquery_test.py index c0ac84732880..458461974ff8 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_test.py @@ -3001,23 +3001,47 @@ def store_callback(table, **kwargs): def test_failure_in_some_rows_does_not_duplicate(self, retry_strategy=None): with mock.patch('time.sleep'): # In this test we simulate a failure to write out two out of three rows. - # Row 0 and row 2 fail to be written on the first attempt, and then - # succeed on the next attempt (if there is one). + # Rows 'value1' and 'value5' fail to be written on the first attempt, and + # then succeed on the next attempt (if there is one). tempdir = '%s%s' % (self._new_tempdir(), os.sep) file_name_1 = os.path.join(tempdir, 'file1_partial') file_name_2 = os.path.join(tempdir, 'file2_partial') + # Errors to return on the first attempt, keyed by 'columnA' of the failed + # rows. Failed rows are identified by their value rather than by their + # position in the batch, since _StreamToBigQuery assigns random shard + # keys and reshuffles the rows, so the order in which they reach the + # client is not guaranteed to match the input order. + first_attempt_errors = { + 'value1': [{ + 'reason': 'i dont like this row' + }, { + 'reason': 'its bad' + }], + 'value5': [{ + 'reason': 'i het this row' + }, { + 'reason': 'its no gud' + }], + } + def store_callback(table, **kwargs): insert_ids = [r for r in kwargs['row_ids']] colA_values = [r['columnA'] for r in kwargs['json_rows']] # The first time this function is called, all rows are included # so we need to filter out 'failed' rows. + passed_indices = [ + i for i, value in enumerate(colA_values) + if value not in first_attempt_errors + ] json_output_1 = { - 'insertIds': [insert_ids[1]], 'colA_values': [colA_values[1]] + 'insertIds': [insert_ids[i] for i in passed_indices], + 'colA_values': [colA_values[i] for i in passed_indices] } - # The second time this function is called, only rows 0 and 2 are incl - # so we don't need to filter any of them. We just write them all out. + # The second time this function is called, only the rows that failed + # the first time are included, so we don't need to filter any of them. + # We just write them all out. json_output_2 = {'insertIds': insert_ids, 'colA_values': colA_values} # The first time we try to insert, we save those insertions in @@ -3025,24 +3049,10 @@ def store_callback(table, **kwargs): if not os.path.exists(file_name_1): with open(file_name_1, 'w') as f: json.dump(json_output_1, f) - return [ - { - 'index': 0, - 'errors': [{ - 'reason': 'i dont like this row' - }, { - 'reason': 'its bad' - }] - }, - { - 'index': 2, - 'errors': [{ - 'reason': 'i het this row' - }, { - 'reason': 'its no gud' - }] - }, - ] + return [{ + 'index': i, 'errors': first_attempt_errors[value] + } for i, value in enumerate(colA_values) + if value in first_attempt_errors] else: with open(file_name_2, 'w') as f: json.dump(json_output_2, f)