Skip to content

[Python] Optimize AvroRowWriter.tell, BigQueryWrapper.insert_rows, and _split_table_spec #41060

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
+24 −9
Draft

[Python] Optimize AvroRowWriter.tell, BigQueryWrapper.insert_rows, and _split_table_spec #41060

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
33 changes: 24 additions & 9 deletions sdks/python/apache_beam/io/gcp/bigquery_tools.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,17 @@
# A single segment of a project id (no dots or colons).
_PROJECT_NAME_SEGMENT_PATTERN = r'[-a-z0-9]*[a-z0-9]'

try:
_COMPILED_TABLE_SPEC_RE = regex.compile(
f'((?P<project>{_PROJECT_PATTERN})[:\\.])?'
f'(?P<dataset>{_DATASET_PATTERN})\\.(?P<table>{_TABLE_PATTERN})')
_COMPILED_PROJECT_RE = regex.compile(_PROJECT_PATTERN)
_COMPILED_PROJECT_SEGMENT_RE = regex.compile(_PROJECT_NAME_SEGMENT_PATTERN)
except NameError:
_COMPILED_TABLE_SPEC_RE = None
_COMPILED_PROJECT_RE = None
_COMPILED_PROJECT_SEGMENT_RE = None

# For CI: temp dataset of this name are automatically deleted
_TEMP_DATASET_PREFIX = 'beam_temp_dataset_'

Expand Down Expand Up @@ -404,10 +415,7 @@ def _split_table_spec(table_spec):
and domain-scoped project ids ('example.com:proj') are both handled. This
mirrors BigQueryHelpers.parseTableSpec in the Java SDK.
"""
pattern = (
f'((?P<project>{_PROJECT_PATTERN})[:\\.])?'
f'(?P<dataset>{_DATASET_PATTERN})\\.(?P<table>{_TABLE_PATTERN})')
if not regex.fullmatch(pattern, table_spec):
if not _COMPILED_TABLE_SPEC_RE.fullmatch(table_spec):
raise _invalid_table_spec(table_spec)

# Table ids cannot contain '.', so the table is always the last segment.
Expand All @@ -426,7 +434,7 @@ def _split_table_spec(table_spec):
return None, prefix, table
project = prefix[:first_dot]
dataset = prefix[first_dot + 1:]
if not regex.fullmatch(_PROJECT_PATTERN, project) or not dataset:
if not _COMPILED_PROJECT_RE.fullmatch(project) or not dataset:
raise _invalid_table_spec(table_spec)
return project, dataset, table

Expand All @@ -440,7 +448,7 @@ def _split_table_spec(table_spec):
dataset = prefix[colon + 1:]
first_dot = dataset.find('.')
if (first_dot >= 0 and first_dot < len(dataset) - 1 and '.' in project and
regex.fullmatch(_PROJECT_NAME_SEGMENT_PATTERN, dataset[:first_dot])):
_COMPILED_PROJECT_SEGMENT_RE.fullmatch(dataset[:first_dot])):
project = project + ':' + dataset[:first_dot]
dataset = dataset[first_dot + 1:]
return project, dataset, table
Expand Down Expand Up @@ -1562,7 +1570,11 @@ def insert_rows(
str(self.unique_row_id) if not insert_ids else insert_ids[i]
for i, _ in enumerate(rows)
]
_json_safe = (str, int, bool, type(None))
rows = [
r if (
type(r) is dict and
all(type(v) in _json_safe for v in r.values())) else

@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: When orjson is installed (the default for fast_json_dumps), integers outside the 64-bit range invoke default_encoder (or raise TypeError), whereas bypassing fast_json_dumps passes arbitrary-precision Python int objects directly to gcp_bigquery.Client.insert_rows_json (which uses standard json.dumps and serializes >64-bit unquoted integers in the JSON payload). Additionally, if r has non-str keys, json.dumps / orjson.dumps behaves differently than passing r through unchanged. Consider verifying type(k) is str as well, or ensuring keys are strings:

all(type(k) is str and type(v) in _json_safe for k, v in r.items())
Suggested change
all(type(v) in _json_safe for v in r.values())) else
all(type(k) is str and type(v) in _json_safe for k, v in r.items())) else

fast_json_loads(fast_json_dumps(r, default=default_encoder))
for r in rows
]
Expand Down Expand Up @@ -1742,7 +1754,7 @@ def encode(self, table_row):
'%s. %s. Row: %r' % (e, JSON_COMPLIANCE_ERROR, table_row))

def decode(self, encoded_table_row):
return json.loads(encoded_table_row.decode('utf-8'))
return json.loads(encoded_table_row)

def to_type_hint(self):
return Any
Expand Down Expand Up @@ -1828,8 +1840,11 @@ def read(self, size=-1):
raise io.UnsupportedOperation("AvroRowWriter is not readable")

def tell(self):
# Flush the fastavro Writer to the underlying stream, otherwise there isn't
# a reliable way to determine how many bytes have been written.
# Include both bytes already flushed to the underlying stream and bytes
# currently buffered in the fastavro Writer, avoiding a costly block flush
# (and 16-byte sync marker) after every single row.
if hasattr(self._avro_writer, 'io'):

@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.

🟢 Low: fastavro.write.Writer in both the Cython (_write_py.pyx) and pure-Python (_write_py.py) implementations exposes .io (a MemoryIO or BytesIO buffer). However, AvroRowWriter.tell() previously had no direct unit test verifying that multiple writes stay buffered within a single block rather than flushing after every tell() call. Consider adding a test in TestAvroRowWriter that calls tell() between write() calls and verifies tell() increases monotonically without forcing separate Avro blocks.

return self._file_handle.tell() + self._avro_writer.io.tell()
self._avro_writer.flush()
return self._file_handle.tell()

Expand Down
Loading