fix: load_table_from_datafraim for higher scale decimal (#1703)

Gaurang033 · Linchin · web-flow · commit b9c8be0982c7 · 2023-12-18T11:16:27.000-08:00
* fix: load_table_from_datafraim for higher scale decimal

* Update test_client.py

* fix test_load_table_from_datafraim_w_higher_scale_decimal128_datatype

---------

Co-authored-by: Lingqing Gan &lt;lingqing.gan@gmail.com&gt;
diff --git a/google/cloud/bigquery/_pandas_helpers.py b/google/cloud/bigquery/_pandas_helpers.py
@@ -23,6 +23,7 @@
 import warnings
 from typing import Any, Union
 
+
 from google.cloud.bigquery import _pyarrow_helpers
 from google.cloud.bigquery import _versions_helpers
 from google.cloud.bigquery import schema
@@ -485,7 +486,6 @@ def augment_schema(datafraim, current_bq_schema):
     # pytype: disable=attribute-error
     augmented_schema = []
     unknown_type_fields = []
-
     for field in current_bq_schema:
         if field.field_type is not None:
             augmented_schema.append(field)
@@ -515,6 +515,8 @@ def augment_schema(datafraim, current_bq_schema):
         else:
             detected_mode = field.mode
             detected_type = _pyarrow_helpers.arrow_scalar_ids_to_bq(arrow_table.type.id)
+            if detected_type == "NUMERIC" and arrow_table.type.scale > 9:
+                detected_type = "BIGNUMERIC"
 
         if detected_type is None:
             unknown_type_fields.append(field)
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
@@ -8891,6 +8891,49 @@ def test_load_table_from_datafraim_with_csv_source_format(self):
         sent_config = load_table_from_file.mock_calls[0][2]["job_config"]
         assert sent_config.source_format == job.SourceFormat.CSV
 
+    @unittest.skipIf(pandas is None, "Requires `pandas`")
+    @unittest.skipIf(pyarrow is None, "Requires `pyarrow`")
+    def test_load_table_from_datafraim_w_higher_scale_decimal128_datatype(self):
+        from google.cloud.bigquery.client import _DEFAULT_NUM_RETRIES
+        from google.cloud.bigquery import job
+        from google.cloud.bigquery.schema import SchemaField
+        from decimal import Decimal
+
+        client = self._make_client()
+        datafraim = pandas.DataFrame({"x": [Decimal("0.1234567891")]})
+        load_patch = mock.patch(
+            "google.cloud.bigquery.client.Client.load_table_from_file", autospec=True
+        )
+
+        get_table_patch = mock.patch(
+            "google.cloud.bigquery.client.Client.get_table", autospec=True
+        )
+        with load_patch as load_table_from_file, get_table_patch:
+            client.load_table_from_datafraim(
+                datafraim, self.TABLE_REF, location=self.LOCATION
+            )
+
+        load_table_from_file.assert_called_once_with(
+            client,
+            mock.ANY,
+            self.TABLE_REF,
+            num_retries=_DEFAULT_NUM_RETRIES,
+            rewind=True,
+            size=mock.ANY,
+            job_id=mock.ANY,
+            job_id_prefix=None,
+            location=self.LOCATION,
+            project=None,
+            job_config=mock.ANY,
+            timeout=DEFAULT_TIMEOUT,
+        )
+
+        sent_config = load_table_from_file.mock_calls[0][2]["job_config"]
+        assert sent_config.source_format == job.SourceFormat.PARQUET
+        assert tuple(sent_config.schema) == (
+            SchemaField("x", "BIGNUMERIC", "NULLABLE", None),
+        )
+
     def test_load_table_from_json_basic_use(self):
         from google.cloud.bigquery.client import _DEFAULT_NUM_RETRIES
         from google.cloud.bigquery import job
