From 4b32a1e852fcbba9e4b14c4cdaa4c05e8698323a Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Tue, 12 Aug 2025 22:02:24 +0000 Subject: [PATCH] perf: Faster session startup by defering anon dataset fetch --- bigframes/session/anonymous_dataset.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/bigframes/session/anonymous_dataset.py b/bigframes/session/anonymous_dataset.py index bc785f693f..ec624d4eb4 100644 --- a/bigframes/session/anonymous_dataset.py +++ b/bigframes/session/anonymous_dataset.py @@ -13,6 +13,7 @@ # limitations under the License. import datetime +import threading from typing import List, Optional, Sequence import uuid @@ -40,19 +41,30 @@ def __init__( ): self.bqclient = bqclient self._location = location - self.dataset = bf_io_bigquery.create_bq_dataset_reference( - self.bqclient, - location=self._location, - ) self.session_id = session_id self._table_ids: List[bigquery.TableReference] = [] self._kms_key = kms_key + self._dataset_lock = threading.Lock() + self._datset_ref: Optional[bigquery.DatasetReference] = None + @property def location(self): return self._location + @property + def dataset(self) -> bigquery.DatasetReference: + if self._datset_ref is not None: + return self._datset_ref + with self._dataset_lock: + if self._datset_ref is None: + self._datset_ref = bf_io_bigquery.create_bq_dataset_reference( + self.bqclient, + location=self._location, + ) + return self._datset_ref + def _default_expiration(self): """When should the table expire automatically?""" return (