Harden multi-row insert casting to deterministic typed behavior.

Replace runtime-value adaptive logic with deterministic casting for SQLAlchemy multi-row VALUES bind markers only, and update compiler tests to verify multi-row casts and single-row non-cast behavior.

Author: msrathore-db

src/databricks/sqlalchemy/_ddl.py

@@ -1,5 +1,4 @@
 import re
-from numbers import Number
 from sqlalchemy.sql import compiler, sqltypes
 import logging
 
@@ -166,21 +165,6 @@ def bindparam_string(self, name, **kw):
             return self._BIND_TEMPLATE % {"name": name.replace("`", "``")}
         return super().bindparam_string(name, **kw)
 
-    @staticmethod
-    def _value_family(value):
-        """Return a coarse runtime family for adaptive multi-row cast decisions."""
-        if value is None:
-            return "null"
-        if isinstance(value, bool):
-            return "bool"
-        if isinstance(value, str):
-            return "string"
-        if isinstance(value, (bytes, bytearray, memoryview)):
-            return "binary"
-        if isinstance(value, Number):
-            return "number"
-        return "other"
-
     @staticmethod
     def _split_multivalue_bind_name(bind_name):
         """Split SQLAlchemy's ``<col>_m<idx>`` bind names into (column, idx)."""
@@ -189,55 +173,37 @@ def _split_multivalue_bind_name(bind_name):
             return None
         return match.group("col"), int(match.group("idx"))
 
-    def _build_adaptive_cast_plan(self):
-        """Return {bind_name: cast_sql_type} for risky multi-row value groups.
+    def _build_multi_value_cast_plan(self, insert_stmt):
+        """Return {bind_name: cast_sql_type} for multi-row VALUES insert binds.
 
-        We only target SQLAlchemy-generated multi-row binds (``*_mN``). For
-        each logical column we inspect row values available at compile time and
-        cast only when families are heterogeneous in a way that commonly causes
-        Spark inline-table incompatibility (e.g., number + string).
+        This is a deterministic fix for Spark inline-table type reconciliation:
+        for SQLAlchemy-generated multi-row INSERT binds (``*_mN``), always cast
+        the marker to the bind's dialect SQL type so each column position in the
+        VALUES table has an explicit server-side type.
         """
-        column_bind_names = {}
-        for bind_name, bind_param in self.binds.items():
-            split = self._split_multivalue_bind_name(bind_name)
-            if split is None:
-                continue
-            column_name, _ = split
-            column_bind_names.setdefault(column_name, []).append((bind_name, bind_param))
+        if not getattr(insert_stmt, "_multi_values", None):
+            return {}
 
         cast_plan = {}
-        for bind_entries in column_bind_names.values():
-            families = set()
-            for _, bind_param in bind_entries:
-                value = getattr(bind_param, "value", None)
-                family = self._value_family(value)
-                if family != "null":
-                    families.add(family)
-
-            if len(families) <= 1:
+        for bind_name, bind_param in self.binds.items():
+            if self._split_multivalue_bind_name(bind_name) is None:
                 continue
 
-            # Numeric + numeric is safe for Spark inline tables and does not
-            # need explicit casting.
-            if families == {"number"}:
+            type_engine = getattr(bind_param, "type", None)
+            if type_engine is None or isinstance(type_engine, sqltypes.NullType):
                 continue
 
-            for bind_name, bind_param in bind_entries:
-                type_engine = getattr(bind_param, "type", None)
-                if type_engine is None or isinstance(type_engine, sqltypes.NullType):
-                    continue
-
-                dialect_type = type_engine._unwrapped_dialect_impl(self.dialect)
-                target_type = self.dialect.type_compiler_instance.process(
-                    dialect_type, identifier_preparer=self.preparer
-                )
-                cast_plan[bind_name] = target_type
+            dialect_type = type_engine._unwrapped_dialect_impl(self.dialect)
+            target_type = self.dialect.type_compiler_instance.process(
+                dialect_type, identifier_preparer=self.preparer
+            )
+            cast_plan[bind_name] = target_type
 
         return cast_plan
 
-    def _apply_adaptive_multi_value_casts(self, sql_text):
+    def _apply_multi_value_casts(self, sql_text, insert_stmt):
         """Wrap selected ``:`name``` markers with ``CAST(... AS <type>)``."""
-        cast_plan = self._build_adaptive_cast_plan()
+        cast_plan = self._build_multi_value_cast_plan(insert_stmt)
         if not cast_plan:
             return sql_text
 
@@ -249,7 +215,7 @@ def _apply_adaptive_multi_value_casts(self, sql_text):
 
     def visit_insert(self, insert_stmt, **kw):
         sql_text = super().visit_insert(insert_stmt, **kw)
-        return self._apply_adaptive_multi_value_casts(sql_text)
+        return self._apply_multi_value_casts(sql_text, insert_stmt)
 
     def limit_clause(self, select, **kw):
         """Identical to the default implementation of SQLCompiler.limit_clause except it writes LIMIT ALL instead of LIMIT -1,

tests/test_local/test_ddl.py

@@ -427,8 +427,8 @@ def test_in_clause_expansion_renders_backticked_markers(self):
         assert ":`col-name_1_3`" in expanded.statement
 
 
-class TestAdaptiveMultiRowInsertCasts(DDLTestBase):
-    def test_mixed_runtime_families_in_multi_values_are_cast(self):
+class TestMultiRowInsertCasts(DDLTestBase):
+    def test_multi_values_casts_mixed_type_column(self):
         metadata = MetaData()
         table = Table("t", metadata, Column("name", String()), Column("value", String()))
         stmt = insert(table).values(
@@ -444,31 +444,38 @@ def test_mixed_runtime_families_in_multi_values_are_cast(self):
         assert "CAST(:`value_m0` AS STRING)" in sql
         assert "CAST(:`value_m1` AS STRING)" in sql
         assert "CAST(:`value_m2` AS STRING)" in sql
-        # Name values are already all string/null and should remain untouched.
-        assert "CAST(:`name_m0` AS STRING)" not in sql
-        assert "CAST(:`name_m1` AS STRING)" not in sql
-        assert "CAST(:`name_m2` AS STRING)" not in sql
+        assert "CAST(:`name_m0` AS STRING)" in sql
+        assert "CAST(:`name_m1` AS STRING)" in sql
+        assert "CAST(:`name_m2` AS STRING)" in sql
 
-    def test_homogeneous_multi_values_are_not_cast(self):
+    def test_homogeneous_multi_values_are_cast(self):
         metadata = MetaData()
         table = Table("t", metadata, Column("value", String()))
         stmt = insert(table).values(
             [{"value": "A"}, {"value": "B"}, {"value": "C"}]
         )
 
         sql = str(stmt.compile(bind=self.engine))
-        assert "CAST(:`value_m0` AS STRING)" not in sql
-        assert "CAST(:`value_m1` AS STRING)" not in sql
-        assert "CAST(:`value_m2` AS STRING)" not in sql
+        assert "CAST(:`value_m0` AS STRING)" in sql
+        assert "CAST(:`value_m1` AS STRING)" in sql
+        assert "CAST(:`value_m2` AS STRING)" in sql
 
-    def test_numeric_family_multi_values_are_not_cast(self):
+    def test_numeric_family_multi_values_are_cast(self):
         metadata = MetaData()
         table = Table("t", metadata, Column("score", Numeric()))
         stmt = insert(table).values(
             [{"score": 1}, {"score": 2.5}, {"score": 3}]
         )
 
         sql = str(stmt.compile(bind=self.engine))
-        assert "CAST(:`score_m0` AS DECIMAL)" not in sql
-        assert "CAST(:`score_m1` AS DECIMAL)" not in sql
-        assert "CAST(:`score_m2` AS DECIMAL)" not in sql
+        assert "CAST(:`score_m0` AS DECIMAL)" in sql
+        assert "CAST(:`score_m1` AS DECIMAL)" in sql
+        assert "CAST(:`score_m2` AS DECIMAL)" in sql
+
+    def test_single_row_insert_does_not_render_casts(self):
+        metadata = MetaData()
+        table = Table("t", metadata, Column("value", String()))
+        stmt = insert(table).values({"value": "A"})
+
+        sql = str(stmt.compile(bind=self.engine))
+        assert "CAST(:`value` AS STRING)" not in sql