Cache BigQuery table definitions in Python SDK #39028
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -28,13 +28,15 @@ | |
| # pytype: skip-file | ||
| # pylint: disable=wrong-import-order, wrong-import-position | ||
|
|
||
| import collections | ||
| import datetime | ||
| import decimal | ||
| import io | ||
| import json | ||
| import logging | ||
| import re | ||
| import sys | ||
| import threading | ||
| import time | ||
| import uuid | ||
| from json.decoder import JSONDecodeError | ||
|
|
@@ -358,6 +360,12 @@ class BigQueryWrapper(object): | |
|
|
||
| HISTOGRAM_METRIC_LOGGER = MetricLogger() | ||
|
|
||
| # Shared by wrapper instances within one Python SDK worker process. | ||
| _TABLE_DEFINITION_CACHE_MAX_ENTRIES = 256 | ||
| _TABLE_DEFINITION_CACHE_TTL_SECS = 60 * 60 | ||
|
Collaborator
There was a problem hiding this comment. I think TTL should be 1 sec. See #34135 (comment) |
||
| _table_definition_cache = collections.OrderedDict() | ||
|
Collaborator
There was a problem hiding this comment. Consider using cachetools.TTLCache instead of hand-rolled cache? |
||
| _table_definition_cache_lock = threading.RLock() | ||
|
|
||
| def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None): | ||
| self.client = client or BigQueryWrapper._bigquery_client(PipelineOptions()) | ||
| self.gcp_bq_client = client or gcp_bigquery.Client( | ||
|
|
@@ -394,6 +402,69 @@ def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None): | |
|
|
||
| self.created_temp_dataset = False | ||
|
|
||
| @classmethod | ||
| def _table_definition_cache_key(cls, project_id, dataset_id, table_id): | ||
| return (project_id, dataset_id, table_id) | ||
|
|
||
| @classmethod | ||
| def _get_cached_table_definition(cls, project_id, dataset_id, table_id): | ||
| cache_key = cls._table_definition_cache_key( | ||
| project_id, dataset_id, table_id) | ||
| now = time.monotonic() | ||
| with cls._table_definition_cache_lock: | ||
| cache_entry = cls._table_definition_cache.get(cache_key) | ||
| if cache_entry is None: | ||
| return None | ||
|
|
||
| expires_at, table = cache_entry | ||
| if expires_at <= now: | ||
| cls._table_definition_cache.pop(cache_key, None) | ||
| return None | ||
|
|
||
| cls._table_definition_cache.move_to_end(cache_key) | ||
| return table | ||
|
|
||
| @classmethod | ||
| def _cache_table_definition(cls, project_id, dataset_id, table_id, table): | ||
| table_type = getattr(bigquery, 'Table', None) | ||
| if table_type is None or not isinstance(table, table_type): | ||
| cls._invalidate_table_definition_cache(project_id, dataset_id, table_id) | ||
| return | ||
|
|
||
| cache_key = cls._table_definition_cache_key( | ||
| project_id, dataset_id, table_id) | ||
| expires_at = time.monotonic() + cls._TABLE_DEFINITION_CACHE_TTL_SECS | ||
| with cls._table_definition_cache_lock: | ||
| cls._table_definition_cache[cache_key] = (expires_at, table) | ||
| cls._table_definition_cache.move_to_end(cache_key) | ||
| while (len(cls._table_definition_cache) | ||
| > cls._TABLE_DEFINITION_CACHE_MAX_ENTRIES): | ||
| cls._table_definition_cache.popitem(last=False) | ||
|
|
||
| @classmethod | ||
| def _invalidate_table_definition_cache( | ||
| cls, project_id=None, dataset_id=None, table_id=None): | ||
| with cls._table_definition_cache_lock: | ||
| if (project_id is not None and dataset_id is not None and | ||
| table_id is not None): | ||
| cache_key = cls._table_definition_cache_key( | ||
| project_id, dataset_id, table_id) | ||
| cls._table_definition_cache.pop(cache_key, None) | ||
| return | ||
|
|
||
|
Comment on lines
+448
to
+454
Contributor
There was a problem hiding this comment. When clearing the entire cache (i.e., when with cls._table_definition_cache_lock:
if project_id is None and dataset_id is None and table_id is None:
cls._table_definition_cache.clear()
return
if (project_id is not None and dataset_id is not None and
table_id is not None):
cache_key = cls._table_definition_cache_key(
project_id, dataset_id, table_id)
cls._table_definition_cache.pop(cache_key, None)
return |
||
| keys_to_delete = [ | ||
| key for key in cls._table_definition_cache | ||
| if ((project_id is None or key[0] == project_id) and | ||
| (dataset_id is None or key[1] == dataset_id) and | ||
| (table_id is None or key[2] == table_id)) | ||
| ] | ||
| for key in keys_to_delete: | ||
| cls._table_definition_cache.pop(key, None) | ||
|
|
||
| @classmethod | ||
| def _clear_table_definition_cache(cls): | ||
| cls._invalidate_table_definition_cache() | ||
|
|
||
| @property | ||
| def unique_row_id(self): | ||
| """Returns a unique row ID (str) used to avoid multiple insertions. | ||
|
|
@@ -804,9 +875,15 @@ def get_table(self, project_id, dataset_id, table_id): | |
| Raises: | ||
| HttpError: if lookup failed. | ||
| """ | ||
| cached_table = self._get_cached_table_definition( | ||
| project_id, dataset_id, table_id) | ||
| if cached_table is not None: | ||
| return cached_table | ||
|
|
||
| request = bigquery.BigqueryTablesGetRequest( | ||
| projectId=project_id, datasetId=dataset_id, tableId=table_id) | ||
| response = self.client.tables.Get(request) | ||
| self._cache_table_definition(project_id, dataset_id, table_id, response) | ||
| return response | ||
|
|
||
| def _create_table( | ||
|
|
@@ -833,6 +910,7 @@ def _create_table( | |
| request = bigquery.BigqueryTablesInsertRequest( | ||
| projectId=project_id, datasetId=dataset_id, table=table) | ||
| response = self.client.tables.Insert(request) | ||
| self._cache_table_definition(project_id, dataset_id, table_id, response) | ||
| _LOGGER.debug("Created the table with id %s", table_id) | ||
| # The response is a bigquery.Table instance. | ||
| return response | ||
|
|
@@ -909,9 +987,12 @@ def _delete_table(self, project_id, dataset_id, table_id): | |
| if exn.status_code == 404: | ||
| _LOGGER.warning( | ||
| 'Table %s:%s.%s does not exist', project_id, dataset_id, table_id) | ||
| self._invalidate_table_definition_cache( | ||
| project_id, dataset_id, table_id) | ||
| return | ||
| else: | ||
| raise | ||
| self._invalidate_table_definition_cache(project_id, dataset_id, table_id) | ||
|
|
||
| @retry.with_exponential_backoff( | ||
| num_retries=MAX_RETRIES, | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
A TTL of 1 hour (
60 * 60seconds) is quite long for caching table definitions. In streaming pipelines or long-running jobs where table schemas might be updated dynamically (e.g., adding a nullable column), workers could use stale metadata for up to an hour, leading to schema mismatch errors. Consider reducing the default TTL to a more conservative value (e.g., 5 or 10 minutes) or making it configurable via pipeline options.