to_gbq respects location argument properly (#225)

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

Author: robertlacok

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

@@ -741,12 +741,12 @@ def test_query_response_bytes(self):
         assert self.gbq_connector.sizeof_fmt(1048576) == "1.0 MB"
         assert self.gbq_connector.sizeof_fmt(1048576000) == "1000.0 MB"
         assert self.gbq_connector.sizeof_fmt(1073741824) == "1.0 GB"
-        assert self.gbq_connector.sizeof_fmt(1.099512E12) == "1.0 TB"
-        assert self.gbq_connector.sizeof_fmt(1.125900E15) == "1.0 PB"
-        assert self.gbq_connector.sizeof_fmt(1.152922E18) == "1.0 EB"
-        assert self.gbq_connector.sizeof_fmt(1.180592E21) == "1.0 ZB"
-        assert self.gbq_connector.sizeof_fmt(1.208926E24) == "1.0 YB"
-        assert self.gbq_connector.sizeof_fmt(1.208926E28) == "10000.0 YB"
+        assert self.gbq_connector.sizeof_fmt(1.099512e12) == "1.0 TB"
+        assert self.gbq_connector.sizeof_fmt(1.125900e15) == "1.0 PB"
+        assert self.gbq_connector.sizeof_fmt(1.152922e18) == "1.0 EB"
+        assert self.gbq_connector.sizeof_fmt(1.180592e21) == "1.0 ZB"
+        assert self.gbq_connector.sizeof_fmt(1.208926e24) == "1.0 YB"
+        assert self.gbq_connector.sizeof_fmt(1.208926e28) == "10000.0 YB"
 
     def test_struct(self, project_id):
         query = """SELECT 1 int_field,
@@ -1325,6 +1325,32 @@ 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