diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index 0cfbf6a2a048..ca2ba0fd2cb0 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -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_PATTERN})[:\\.])?' + f'(?P{_DATASET_PATTERN})\\.(?P{_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_' @@ -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_PATTERN})[:\\.])?' - f'(?P{_DATASET_PATTERN})\\.(?P
{_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. @@ -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 @@ -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 @@ -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 fast_json_loads(fast_json_dumps(r, default=default_encoder)) for r in rows ] @@ -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 @@ -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'): + return self._file_handle.tell() + self._avro_writer.io.tell() self._avro_writer.flush() return self._file_handle.tell()