SNOW-1460351: support case when string type coercion

CHANGELOG.md

@@ -22,6 +22,7 @@
 - Fixed a bug in convert_timezone that made the setting the source_timezone parameter return an error.
 - Fixed a bug where creating DataFrame with empty data of type `DateType` raises `AttributeError`.
 - Fixed a bug that table merge fails when update clause exists but no update takes place.
+- Fixed a bug that case when expression did not handle string type coercion.
 
 ### Snowpark pandas API Updates
 

src/snowflake/snowpark/mock/_plan.py

@@ -39,6 +39,7 @@
     UnboundedPreceding,
     WindowExpression,
 )
+from snowflake.snowpark.mock._snowflake_data_type import coerce_type_if_needed
 from snowflake.snowpark.mock._udf_utils import (
     coerce_variant_input,
     remove_null_wrapper,
@@ -1992,13 +1993,21 @@ def _match_pattern(row) -> bool:
             remaining = remaining[~remaining.index.isin(true_index)]
 
             if output_data.sf_type:
-                if (
-                    not isinstance(output_data.sf_type.datatype, NullType)
-                    and output_data.sf_type != value.sf_type
+                if not isinstance(output_data.sf_type.datatype, NullType) and not (
+                    types_are_compatible(
+                        output_data.sf_type.datatype, value.sf_type.datatype
+                    )
                 ):
                     raise SnowparkLocalTestingException(
                         f"CaseWhen expressions have conflicting data types: {output_data.sf_type} != {value.sf_type}"
                     )
+                else:
+                    output_data.sf_type = ColumnType(
+                        coerce_type_if_needed(
+                            output_data.sf_type.datatype, value.sf_type.datatype
+                        ),
+                        output_data.sf_type.nullable,
+                    )
             else:
                 output_data.sf_type = value.sf_type
 
@@ -2008,13 +2017,21 @@ def _match_pattern(row) -> bool:
             )
             output_data[remaining.index] = value[remaining.index]
             if output_data.sf_type:
-                if (
-                    not isinstance(output_data.sf_type.datatype, NullType)
-                    and output_data.sf_type.datatype != value.sf_type.datatype
+                if not isinstance(output_data.sf_type.datatype, NullType) and not (
+                    types_are_compatible(
+                        output_data.sf_type.datatype, value.sf_type.datatype
+                    )
                 ):
                     raise SnowparkLocalTestingException(
                         f"CaseWhen expressions have conflicting data types: {output_data.sf_type.datatype} != {value.sf_type.datatype}"
                     )
+                else:
+                    output_data.sf_type = ColumnType(
+                        coerce_type_if_needed(
+                            output_data.sf_type.datatype, value.sf_type.datatype
+                        ),
+                        output_data.sf_type.nullable,
+                    )
             else:
                 output_data.sf_type = value.sf_type
         return output_data

src/snowflake/snowpark/mock/_snowflake_data_type.py

@@ -15,6 +15,7 @@
     FloatType,
     IntegerType,
     LongType,
+    StringType,
     _IntegralType,
     _NumericType,
 )
@@ -530,3 +531,12 @@ def isnull(self):
         result = super().isnull()
         result.sf_type = ColumnType(BooleanType(), True)
         return result
+
+
+def coerce_type_if_needed(type1: DataType, type2: DataType) -> DataType:
+    from snowflake.snowpark.mock._udf_utils import types_are_compatible
+
+    if types_are_compatible(type1, type2):
+        if isinstance(type1, StringType) and isinstance(type2, StringType):
+            return StringType(max(type1.length, type2.length))
+    return type1

tests/integ/test_datatypes.py

@@ -6,7 +6,7 @@
 import pytest
 
 from snowflake.snowpark import DataFrame, Row
-from snowflake.snowpark.functions import lit
+from snowflake.snowpark.functions import col, lit, when
 from snowflake.snowpark.types import (
     BooleanType,
     DecimalType,
@@ -423,3 +423,12 @@ def test_join_basic(session):
             ]
         )
     )
+
+
+def test_case_when_type_coerce(session):
+    df_input = session.create_dataframe(
+        [1], StructType([StructField("col", LongType())])
+    )
+    assert df_input.withColumn(
+        "new_col", when((col("col").is_null()), lit("abc")).otherwise(lit("abcdef"))
+    ).collect() == [Row(1, "abcdef")]