diff --git a/sdk/python/feast/infra/offline_stores/bigquery.py b/sdk/python/feast/infra/offline_stores/bigquery.py index 15c077ecbff..1dc55c43941 100644 --- a/sdk/python/feast/infra/offline_stores/bigquery.py +++ b/sdk/python/feast/infra/offline_stores/bigquery.py @@ -136,7 +136,10 @@ def get_historical_features( assert isinstance(config.offline_store, BigQueryOfflineStoreConfig) table_reference = _get_table_reference_for_new_entity( - client, client.project, config.offline_store.dataset + client, + client.project, + config.offline_store.dataset, + config.offline_store.location, ) @contextlib.contextmanager @@ -340,13 +343,16 @@ def _wait_until_done(bq_job): def _get_table_reference_for_new_entity( - client: Client, dataset_project: str, dataset_name: str + client: Client, + dataset_project: str, + dataset_name: str, + dataset_location: Optional[str], ) -> str: """Gets the table_id for the new entity to be uploaded.""" # First create the BigQuery dataset if it doesn't exist dataset = bigquery.Dataset(f"{dataset_project}.{dataset_name}") - dataset.location = "US" + dataset.location = dataset_location if dataset_location else "US" try: client.get_dataset(dataset)