[damccorm]

Description

Optimizes three hot paths in sdks/python/apache_beam/io/gcp/bigquery_tools.py:

1. Avoid per-row Avro block flushes in AvroRowWriter.tell():
   • During BigQuery FILE_LOADS with temp_file_format=AVRO, the file writer calls AvroRowWriter.tell() after every row to check shard rollover limits. Previously, AvroRowWriter.tell() called self._avro_writer.flush() before self._file_handle.tell(), forcing fastavro to close, compress, and write a 1-row Avro block (plus a 16-byte sync marker) on every single record.
   • Returning self._file_handle.tell() + self._avro_writer.io.tell() includes buffered bytes without flushing the active Avro block (16.05x faster AvroRowWriter throughput and 12.2x smaller output files on 50,000 rows).
2. Fast-path JSON-safe scalar rows in BigQueryWrapper.insert_rows and RowAsDictJsonCoder.decode:
   • Skips the redundant fast_json_loads(fast_json_dumps(r)) round-trip in BigQueryWrapper.insert_rows when a row dict contains only JSON-safe scalar values (str, int, bool, None) (11.61x faster row preparation), and passes bytes directly to json.loads() in RowAsDictJsonCoder.decode.
3. Precompile table spec regular expressions in _split_table_spec:
   • Hoists _COMPILED_TABLE_SPEC_RE, _COMPILED_PROJECT_RE, and _COMPILED_PROJECT_SEGMENT_RE to module initialization (2.34x faster _split_table_spec).

Testing

• Verified with pytest sdks/python/apache_beam/io/gcp/bigquery_tools_test.py.

[damccorm]

/gemini review

[github-actions]

📋 Review Summary

This PR introduces three targeted performance optimizations in bigquery_tools.py: avoiding per-row Avro block flushes in AvroRowWriter.tell(), fast-pathing JSON-safe scalar rows in BigQueryWrapper.insert_rows() and RowAsDictJsonCoder.decode(), and precompiling regular expressions for _split_table_spec(). The changes are well-reasoned and significantly improve throughput on hot paths.

🔍 General Feedback

• In BigQueryWrapper.insert_rows, the fast-path for int values should guard against integers exceeding 64-bit bounds, which orjson previously routed to default_encoder.
• Adding unit tests for AvroRowWriter.tell() and the insert_rows fast-path/fallback cases in bigquery_tools_test.py would help prevent future regressions.

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

[github-actions]

🟡 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

[github-actions]

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