Create the BQ dataset in the location specified in feature_store.yaml if the dataset doesn't exists (#1959)

francoisserra · achals · web-flow · commit 73af2fa9ca73 · 2021-11-18T11:01:38.000-08:00
* fixes #1958 Signed-off-by: François Serra <francois.serra@adeo.com> * Fix lint Signed-off-by: Achal Shah <achals@gmail.com> * Fix lint Signed-off-by: Achal Shah <achals@gmail.com> Co-authored-by: Achal Shah <achals@gmail.com>
diff --git 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)
