fix: load_table_from_dataframe for higher scale decimal (#1703)

* fix: load_table_from_dataframe for higher scale decimal * Update test_client.py * fix test_load_table_from_dataframe_w_higher_scale_decimal128_datatype --------- Co-authored-by: Lingqing Gan <[email protected]>
googleapis · Dec 18, 2023 · b9c8be0 · b9c8be0
1 parent fdd580f
commit b9c8be0
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 1 deletion.
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(dataframe, 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(dataframe, 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_dataframe_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_dataframe_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()
+ dataframe = 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_dataframe(
+ dataframe, 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