to_gbq respects location argument properly

If dataset does not exist, it gets created in the correct location

Author: Robert Lacok

pandas_gbq/gbq.py

@@ -283,7 +283,7 @@ def __init__(
 
         # BQ Queries costs $5 per TB. First 1 TB per month is free
         # see here for more: https://cloud.google.com/bigquery/pricing
-        self.query_price_for_TB = 5. / 2 ** 40  # USD/TB
+        self.query_price_for_TB = 5.0 / 2 ** 40  # USD/TB
 
     def _start_timer(self):
         self.start = time.time()
@@ -895,7 +895,11 @@ def to_gbq(
     dataset_id, table_id = destination_table.rsplit(".", 1)
 
     table = _Table(
-        project_id, dataset_id, reauth=reauth, private_key=private_key
+        project_id,
+        dataset_id,
+        reauth=reauth,
+        private_key=private_key,
+        location=location,
     )
 
     if not table_schema:
@@ -967,9 +971,18 @@ def _generate_bq_schema(df, default_type="STRING"):
 
 
 class _Table(GbqConnector):
-    def __init__(self, project_id, dataset_id, reauth=False, private_key=None):
+    def __init__(
+        self,
+        project_id,
+        dataset_id,
+        reauth=False,
+        private_key=None,
+        location=None,
+    ):
         self.dataset_id = dataset_id
-        super(_Table, self).__init__(project_id, reauth, private_key)
+        super(_Table, self).__init__(
+            project_id, reauth, private_key, location=location
+        )
 
     def exists(self, table_id):
         """ Check if a table exists in Google BigQuery
@@ -1017,9 +1030,11 @@ def create(self, table_id, schema):
         if not _Dataset(self.project_id, private_key=self.private_key).exists(
             self.dataset_id
         ):
-            _Dataset(self.project_id, private_key=self.private_key).create(
-                self.dataset_id
-            )
+            _Dataset(
+                self.project_id,
+                private_key=self.private_key,
+                location=self.location,
+            ).create(self.dataset_id)
 
         table_ref = self.client.dataset(self.dataset_id).table(table_id)
         table = Table(table_ref)
@@ -1064,8 +1079,12 @@ def delete(self, table_id):
 
 
 class _Dataset(GbqConnector):
-    def __init__(self, project_id, reauth=False, private_key=None):
-        super(_Dataset, self).__init__(project_id, reauth, private_key)
+    def __init__(
+        self, project_id, reauth=False, private_key=None, location=None
+    ):
+        super(_Dataset, self).__init__(
+            project_id, reauth, private_key, location=location
+        )
 
     def exists(self, dataset_id):
         """ Check if a dataset exists in Google BigQuery
@@ -1107,6 +1126,9 @@ def create(self, dataset_id):
 
         dataset = Dataset(self.client.dataset(dataset_id))
 
+        if self.location is not None:
+            dataset.location = self.location
+
         try:
             self.client.create_dataset(dataset)
         except self.http_error as ex:

tests/system/test_gbq.py

@@ -1325,6 +1325,26 @@ def test_upload_data_tokyo(
         )
         assert table.num_rows > 0
 
+    def test_upload_data_tokyo_non_existing_dataset(self, project_id, random_dataset_id, bigquery_client):
+        test_size = 10
+        df = make_mixed_dataframe_v2(test_size)
+        non_existing_tokyo_dataset = random_dataset_id
+        non_existing_tokyo_destination = "{}.to_gbq_test".format(non_existing_tokyo_dataset)
+
+        # Initialize table with sample data
+        gbq.to_gbq(
+            df,
+            non_existing_tokyo_destination,
+            project_id,
+            private_key=self.credentials,
+            location="asia-northeast1",
+        )
+
+        table = bigquery_client.get_table(
+            bigquery_client.dataset(non_existing_tokyo_dataset).table("to_gbq_test")
+        )
+        assert table.num_rows > 0
+
 
 # _Dataset tests