Skip to content

[Python] Optimize BigQueryWriteFn row size estimation and table ref parsing #41059

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
damccorm wants to merge 1 commit into
base: master
Choose a base branch
from
+31 −9
Draft

[Python] Optimize BigQueryWriteFn row size estimation and table ref parsing #41059

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 31 additions & 9 deletions sdks/python/apache_beam/io/gcp/bigquery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -366,6 +366,7 @@ def chain_after(result):
import logging
import random
import secrets
import sys
import time
import uuid
import warnings
Expand Down Expand Up @@ -1522,6 +1523,25 @@ def BigQuerySink(*args, validate=False, **kwargs):


_KNOWN_TABLES = set()
_PRIMITIVE_SIZEOF_TYPES = (str, int, float, bool, bytes, type(None))


def _fast_row_and_insert_id_size(row_and_insert_id):
if type(row_and_insert_id) is tuple and len(row_and_insert_id) == 2:
row, insert_id = row_and_insert_id
if type(row) is dict and type(insert_id) in _PRIMITIVE_SIZEOF_TYPES:
getsizeof = sys.getsizeof
total = getsizeof(row_and_insert_id) + getsizeof(row)
if insert_id is not None:
total += getsizeof(insert_id)
for k, v in row.items():
if not (type(k) in _PRIMITIVE_SIZEOF_TYPES and
type(v) in _PRIMITIVE_SIZEOF_TYPES):
return get_deep_size(row_and_insert_id)
total += getsizeof(k) + (

@github-actions github-actions Bot Nov 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Medium: objsize.get_deep_size tracks visited object IDs so that objects referenced multiple times within the same row (e.g., multiple columns referencing the same string/bytes object, or cached/interned scalars like True/False, small ints, or "") are only counted once. Here, v is only deduplicated against insert_id, so if multiple keys in row share the same value reference (especially a large str or bytes object), total will count sys.getsizeof(v) multiple times. This could cause a row to be overestimated relative to get_deep_size and falsely exceed _max_insert_payload_size (routing valid rows to FAILED_ROWS).

Consider either falling back to get_deep_size when total >= max_insert_payload_size in _check_row_size before rejecting a row, or tracking seen object IDs when needed.

getsizeof(v) if v is not None and v is not insert_id else 0)
return total
return get_deep_size(row_and_insert_id)


class BigQueryWriteFn(DoFn):
Expand Down Expand Up @@ -1727,7 +1747,7 @@ def _create_table_if_needed(self, table_reference, schema=None):

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)
row_byte_size = _fast_row_and_insert_id_size(row_and_insert_id)

@github-actions github-actions Bot Nov 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Medium: If _fast_row_and_insert_id_size overestimates a row due to shared object references across fields (which get_deep_size deduplicates), a row near the limit could be routed to the DLQ even though get_deep_size(row_and_insert_id) < self._max_insert_payload_size. Re-checking with get_deep_size(row_and_insert_id) inside the if row_byte_size >= self._max_insert_payload_size: branch would preserve exact backwards-compatible DLQ behavior at zero cost on the normal path.


# Check if individual row exceeds size limit
if row_byte_size >= self._max_insert_payload_size:
Expand All @@ -1744,15 +1764,17 @@ def process(
self, element, window_value=DoFn.WindowedValueParam, *schema_side_inputs):
destination = bigquery_tools.get_hashable_destination(element[0])

if callable(self.schema):
schema = self.schema(destination, *schema_side_inputs)
elif isinstance(self.schema, vp.ValueProvider):
schema = self.schema.get()
else:
schema = self.schema
if (self.create_disposition != BigQueryDisposition.CREATE_NEVER and
destination not in _KNOWN_TABLES):
if callable(self.schema):
schema = self.schema(destination, *schema_side_inputs)
elif isinstance(self.schema, vp.ValueProvider):
schema = self.schema.get()
else:
schema = self.schema

self._create_table_if_needed(
bigquery_tools.parse_table_reference(destination), schema)
self._create_table_if_needed(
bigquery_tools.parse_table_reference(destination), schema)

if not self.with_batched_input:
row_and_insert_id = element[1]
Expand Down
Loading