estimate request size in BQ streaming inserts when with_auto_sharding is true (#35212)

Author: quentin-sommer

sdks/python/apache_beam/io/gcp/bigquery.py

@@ -1590,6 +1590,21 @@ def _create_table_if_needed(self, table_reference, schema=None):
         additional_create_parameters=self.additional_bq_parameters)
     _KNOWN_TABLES.add(str_table_reference)
 
+  def _check_row_size(self, row_and_insert_id) -> Tuple[int, Optional[str]]:
+    """Returns error string when the row estimated size is too big"""
+    row_byte_size = get_deep_size(row_and_insert_id)
+
+    # Check if individual row exceeds size limit
+    if row_byte_size >= self._max_insert_payload_size:
+      row_mb_size = row_byte_size / 1_000_000
+      max_mb_size = self._max_insert_payload_size / 1_000_000
+      return (
+          row_byte_size,
+          (
+              f"Received row with size {row_mb_size}MB that exceeds "
+              f"the maximum insert payload size set ({max_mb_size}MB)."))
+    return (row_byte_size, None)
+
   def process(
       self, element, window_value=DoFn.WindowedValueParam, *schema_side_inputs):
     destination = bigquery_tools.get_hashable_destination(element[0])
@@ -1606,20 +1621,15 @@ def process(
 
     if not self.with_batched_input:
       row_and_insert_id = element[1]
-      row_byte_size = get_deep_size(row_and_insert_id)
+      row_byte_size, row_too_big_error = self._check_row_size(row_and_insert_id)
 
       # send large rows that exceed BigQuery insert limits to DLQ
-      if row_byte_size >= self._max_insert_payload_size:
-        row_mb_size = row_byte_size / 1_000_000
-        max_mb_size = self._max_insert_payload_size / 1_000_000
-        error = (
-            f"Received row with size {row_mb_size}MB that exceeds "
-            f"the maximum insert payload size set ({max_mb_size}MB).")
+      if row_too_big_error is not None:
         return [
             pvalue.TaggedOutput(
                 BigQueryWriteFn.FAILED_ROWS_WITH_ERRORS,
                 window_value.with_value(
-                    (destination, row_and_insert_id[0], error))),
+                    (destination, row_and_insert_id[0], row_too_big_error))),
             pvalue.TaggedOutput(
                 BigQueryWriteFn.FAILED_ROWS,
                 window_value.with_value((destination, row_and_insert_id[0])))
@@ -1643,11 +1653,50 @@ def process(
       if self._total_buffered_rows >= self._max_buffered_rows:
         return self._flush_all_batches()
     else:
-      # The input is already batched per destination, flush the rows now.
+      # The input is already batched per destination
+      # but we verify the payload size.
+      # The batch might be split into smaller batches
       batched_rows = element[1]
-      for r in batched_rows:
-        self._rows_buffer[destination].append((r, window_value))
-      return self._flush_batch(destination)
+      failed_outputs = []
+      current_batch = []
+      current_batch_size = 0
+
+      for row in batched_rows:
+        row_byte_size, row_too_big_error = self._check_row_size(row)
+        # Check if individual row exceeds size limit
+        if row_too_big_error is not None:
+          failed_outputs.extend([
+              pvalue.TaggedOutput(
+                  BigQueryWriteFn.FAILED_ROWS_WITH_ERRORS,
+                  window_value.with_value(
+                      (destination, row[0], row_too_big_error))),
+              pvalue.TaggedOutput(
+                  BigQueryWriteFn.FAILED_ROWS,
+                  window_value.with_value((destination, row[0])))
+          ])
+          continue
+
+        # Check if adding this row would exceed batch size limit
+        if (len(current_batch) != 0 and
+            current_batch_size + row_byte_size > self._max_insert_payload_size):
+
+          self._rows_buffer[destination].extend(
+              ((batch_row, window_value) for batch_row in current_batch))
+          failed_outputs.extend(self._flush_batch(destination))
+
+          # Start new batch with current row
+          current_batch = [row]
+          current_batch_size = row_byte_size
+        else:
+          current_batch.append(row)
+          current_batch_size += row_byte_size
+
+      if current_batch:
+        for batch_row in current_batch:
+          self._rows_buffer[destination].append((batch_row, window_value))
+        failed_outputs.extend(self._flush_batch(destination))
+
+      return failed_outputs
 
   def finish_bundle(self):
     bigquery_tools.BigQueryWrapper.HISTOGRAM_METRIC_LOGGER.log_metrics(

sdks/python/apache_beam/io/gcp/bigquery_test.py

@@ -1944,6 +1944,147 @@ def store_callback(table, **kwargs):
       assert_that(failed_values, equal_to(failed_rows))
 
 
+def test_with_batched_input_exceeds_size_limit(self):
+  from apache_beam.utils.windowed_value import WindowedValue
+  from apache_beam.transforms import window
+
+  client = mock.Mock()
+  client.tables.Get.return_value = bigquery.Table(
+      tableReference=bigquery.TableReference(
+          projectId='project-id', datasetId='dataset_id', tableId='table_id'))
+  client.insert_rows_json.return_value = []
+  fn = beam.io.gcp.bigquery.BigQueryWriteFn(
+      batch_size=10,
+      create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
+      write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND,
+      kms_key=None,
+      with_batched_input=True,
+      max_insert_payload_size=500,  # Small size to trigger the check
+      test_client=client)
+
+  fn.start_bundle()
+
+  # Create rows where the third one exceeds the size limit
+  rows = [
+      ({
+          'month': 1
+      }, 'insertid1'),
+      ({
+          'month': 2
+      }, 'insertid2'),
+      ({
+          'columnA': 'large_string' * 100
+      }, 'insertid3')  # This exceeds 500 bytes
+  ]
+
+  # Create a windowed value for testing
+  test_window = window.GlobalWindow()
+  test_timestamp = window.MIN_TIMESTAMP
+  windowed_value = WindowedValue(None, test_timestamp, [test_window])
+
+  # Process the batched input
+  result = fn.process(('project-id:dataset_id.table_id', rows), windowed_value)
+
+  # Convert generator to list to check results
+  failed_rows = list(result)
+
+  # Should have 2 failed outputs (FAILED_ROWS_WITH_ERRORS and FAILED_ROWS)
+  self.assertEqual(len(failed_rows), 2)
+
+  # Check that the large row was sent to DLQ
+  failed_with_errors = [
+      r for r in failed_rows if r.tag == fn.FAILED_ROWS_WITH_ERRORS
+  ]
+  failed_without_errors = [r for r in failed_rows if r.tag == fn.FAILED_ROWS]
+
+  self.assertEqual(len(failed_with_errors), 1)
+  self.assertEqual(len(failed_without_errors), 1)
+
+  # Verify the error message
+  destination, row, error = failed_with_errors[0].value.value
+  self.assertEqual(destination, 'project-id:dataset_id.table_id')
+  self.assertEqual(row, {'columnA': 'large_string' * 100})
+  self.assertIn('exceeds the maximum insert payload size', error)
+
+  # Verify that only the valid rows were sent to BigQuery
+  self.assertTrue(client.insert_rows_json.called)
+  # Check that the rows were inserted (might be in multiple batches)
+  total_rows = []
+  for call in client.insert_rows_json.call_args_list:
+    total_rows.extend(call[1]['json_rows'])
+
+  self.assertEqual(len(total_rows), 2)
+  self.assertEqual(total_rows[0], {'month': 1})
+  self.assertEqual(total_rows[1], {'month': 2})
+
+
+def test_with_batched_input_splits_large_batch(self):
+  from apache_beam.utils.windowed_value import WindowedValue
+  from apache_beam.transforms import window
+
+  client = mock.Mock()
+  client.tables.Get.return_value = bigquery.Table(
+      tableReference=bigquery.TableReference(
+          projectId='project-id', datasetId='dataset_id', tableId='table_id'))
+  client.insert_rows_json.return_value = []
+  create_disposition = beam.io.BigQueryDisposition.CREATE_IF_NEEDED
+  write_disposition = beam.io.BigQueryDisposition.WRITE_APPEND
+
+  fn = beam.io.gcp.bigquery.BigQueryWriteFn(
+      batch_size=10,
+      create_disposition=create_disposition,
+      write_disposition=write_disposition,
+      kms_key=None,
+      with_batched_input=True,
+      max_insert_payload_size=800,  # Small size to force batch splitting
+      test_client=client)
+
+  fn.start_bundle()
+
+  # Create rows that together exceed the batch size limit.
+  # Each row with 'data' * 10 is about 200 bytes
+  # So 2 rows should fit, 3rd should cause flush.
+  rows = [
+      ({
+          'data': 'x' * 10
+      }, 'insertid1'),
+      ({
+          'data': 'y' * 10
+      }, 'insertid2'),
+      ({
+          'data': 'z' * 10
+      }, 'insertid3'),  # This should go in second batch
+  ]
+
+  # Create a windowed value for testing
+  test_window = window.GlobalWindow()
+  test_timestamp = window.MIN_TIMESTAMP
+  windowed_value = WindowedValue(None, test_timestamp, [test_window])
+
+  # Process the batched input
+  result = fn.process(('project-id:dataset_id.table_id', rows), windowed_value)
+
+  # Convert generator to list (should be empty as no failures)
+  failed_rows = list(result)
+  self.assertEqual(len(failed_rows), 0)
+
+  # With 800 byte limit and ~341 byte rows
+  # we should be able to fit 2 rows per batch.
+  # So we expect 2 calls: [row1, row2] and [row3]
+  self.assertEqual(client.insert_rows_json.call_count, 2)
+
+  # Check first batch (2 rows)
+  first_call = client.insert_rows_json.call_args_list[0][1]
+  self.assertEqual(len(first_call['json_rows']), 2)
+  self.assertEqual(first_call['json_rows'][0], {'data': 'x' * 10})
+  self.assertEqual(first_call['json_rows'][1], {'data': 'y' * 10})
+
+  # Check second batch (1 row)
+  second_call = client.insert_rows_json.call_args_list[1][1]
+  self.assertEqual(len(second_call['json_rows']), 1)
+  self.assertEqual(second_call['json_rows'][0], {'data': 'z' * 10})
+
+
 @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
 class BigQueryStreamingInsertTransformTests(unittest.TestCase):
   def test_dofn_client_process_performs_batching(self):