Fix: Support coercion of BIGNUMERIC values into FLOAT64 for BigQuery (#3417)

Author: izeigerman

sqlmesh/core/engine_adapter/bigquery.py

@@ -78,6 +78,11 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, Row
                 exp.DataType.build("DATETIME", dialect=DIALECT),
             },
         },
+        coerceable_types={
+            exp.DataType.build("FLOAT64", dialect=DIALECT): {
+                exp.DataType.build("BIGNUMERIC", dialect=DIALECT),
+            },
+        },
         support_coercing_compatible_types=True,
         parameterized_type_defaults={
             exp.DataType.build("DECIMAL", dialect=DIALECT).this: [(38, 9), (0,)],

sqlmesh/core/schema_diff.py

@@ -2,8 +2,9 @@
 
 import logging
 import typing as t
-from collections import defaultdict
 from enum import Enum, auto
+from collections import defaultdict
+from pydantic import Field
 from sqlglot import exp
 from sqlglot.helper import ensure_list, seq_get
 
@@ -305,6 +306,11 @@ class SchemaDiffer(PydanticModel):
             columns of nested STRUCTs.
         compatible_types: Types that are compatible and automatically coerced in actions like UNION ALL. Dict key is data
             type, and value is the set of types that are compatible with it.
+        coerceable_types: The mapping from a current type to all types that can be safely coerced to the current one without
+            altering the column type. NOTE: usually callers should not specify this attribute manually and set the
+            `support_coercing_compatible_types` flag instead. Some engines are inconsistent about their type coercion rules.
+            For example, in BigQuery a BIGNUMERIC column can't be altered to be FLOAT64, while BIGNUMERIC values can be inserted
+            into a FLOAT64 column just fine.
         support_coercing_compatible_types: Whether or not the engine for which the diff is being computed supports direct
             coercion of compatible types.
         parameterized_type_defaults: Default values for parameterized data types. Dict key is a sqlglot exp.DataType.Type,
@@ -326,6 +332,9 @@ class SchemaDiffer(PydanticModel):
     support_nested_drop: bool = False
     array_element_selector: str = ""
     compatible_types: t.Dict[exp.DataType, t.Set[exp.DataType]] = {}
+    coerceable_types_: t.Dict[exp.DataType, t.Set[exp.DataType]] = Field(
+        default_factory=dict, alias="coerceable_types"
+    )
     support_coercing_compatible_types: bool = False
     parameterized_type_defaults: t.Dict[
         exp.DataType.Type, t.List[t.Tuple[t.Union[int, float], ...]]
@@ -339,8 +348,9 @@ class SchemaDiffer(PydanticModel):
     def coerceable_types(self) -> t.Dict[exp.DataType, t.Set[exp.DataType]]:
         if not self._coerceable_types:
             if not self.support_coercing_compatible_types or not self.compatible_types:
-                return {}
-            coerceable_types = defaultdict(set)
+                return self.coerceable_types_
+            coerceable_types: t.Dict[exp.DataType, t.Set[exp.DataType]] = defaultdict(set)
+            coerceable_types.update(self.coerceable_types_)
             for source_type, target_types in self.compatible_types.items():
                 for target_type in target_types:
                     coerceable_types[target_type].add(source_type)
@@ -361,8 +371,6 @@ def _is_compatible_type(self, current_type: exp.DataType, new_type: exp.DataType
         return False
 
     def _is_coerceable_type(self, current_type: exp.DataType, new_type: exp.DataType) -> bool:
-        if not self.support_coercing_compatible_types:
-            return False
         if current_type in self.coerceable_types:
             is_coerceable = new_type in self.coerceable_types[current_type]
             if is_coerceable:

tests/core/engine_adapter/test_base.py

@@ -563,6 +563,26 @@ def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture)
                 'ALTER TABLE "test_table" ADD COLUMN "array_col"."element"."array_b" INT AFTER "array_a"',
             ],
         ),
+        (
+            {
+                "coerceable_types": {
+                    exp.DataType.build("FLOAT"): {exp.DataType.build("INT")},
+                },
+            },
+            {
+                "a": "FLOAT",
+                "b": "TEXT",
+            },
+            {
+                "a": "INT",
+                "b": "TEXT",
+            },
+            {
+                "a": "FLOAT",
+                "b": "TEXT",
+            },
+            [],
+        ),
         (
             {
                 "support_nested_operations": True,
@@ -865,6 +885,8 @@ def _from_structs(
         current_struct: exp.DataType, new_struct: exp.DataType
     ) -> t.List[TableAlterOperation]:
         operations = original_from_structs(current_struct, new_struct)
+        if not operations:
+            return operations
         assert (
             operations[-1].expected_table_struct.sql()
             == columns_to_types_to_struct(expected_final_structure).sql()

tests/core/test_schema_diff.py

@@ -1108,6 +1108,34 @@ def test_schema_diff_calculate_type_transitions():
                 },
             ),
         ),
+        (
+            "STRUCT<id INT, name STRING, revenue FLOAT>",
+            "STRUCT<id INT, name STRING, revenue INT>",
+            [],
+            dict(
+                support_positional_add=True,
+                support_nested_operations=True,
+                coerceable_types={
+                    exp.DataType.build("FLOAT"): {exp.DataType.build("INT")},
+                },
+            ),
+        ),
+        (
+            "STRUCT<id INT, name STRING, revenue FLOAT>",
+            "STRUCT<id INT, name INT, revenue INT>",
+            [],
+            dict(
+                support_positional_add=True,
+                support_nested_operations=True,
+                support_coercing_compatible_types=True,
+                compatible_types={
+                    exp.DataType.build("INT"): {exp.DataType.build("FLOAT")},
+                },
+                coerceable_types={
+                    exp.DataType.build("STRING"): {exp.DataType.build("INT")},
+                },
+            ),
+        ),
         # Coercion with an alter results in a single alter
         (
             "STRUCT<id INT, name STRING, revenue FLOAT, total INT>",