snowpark-connect 0.32.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

snowflake/snowpark_connect/expression/map_unresolved_function.py

@@ -24,7 +24,6 @@ from typing import List, Optional
 from urllib.parse import quote, unquote
 
 import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
-import pyspark.sql.functions as pyspark_functions
 from google.protobuf.message import Message
 from pyspark.errors.exceptions.base import (
     AnalysisException,
@@ -101,6 +100,7 @@ from snowflake.snowpark_connect.expression.map_unresolved_star import (
 )
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.relation.catalogs.utils import CURRENT_CATALOG_NAME
+from snowflake.snowpark_connect.relation.utils import is_aggregate_function
 from snowflake.snowpark_connect.type_mapping import (
     map_json_schema_to_snowpark,
     map_pyspark_types_to_snowpark_types,
@@ -119,7 +119,6 @@ from snowflake.snowpark_connect.utils.context import (
     get_is_evaluating_sql,
     get_is_in_udtf_context,
     get_spark_version,
-    is_in_pivot,
     is_window_enabled,
     push_udtf_context,
     resolving_fun_args,
@@ -400,9 +399,8 @@ def map_unresolved_function(
     result_type: Optional[DataType | List[DateType]] = None
     qualifier_parts: List[str] = []
 
-    pyspark_func = getattr(pyspark_functions, function_name, None)
-    if pyspark_func and pyspark_func.__doc__.lstrip().startswith("Aggregate function:"):
-        # Used by the GROUP BY ALL implementation. Far from ideal, but it seems to work...
+    # Check if this is an aggregate function (used by GROUP BY ALL implementation)
+    if is_aggregate_function(function_name):
         add_sql_aggregate_function()
 
     def _type_with_typer(col: Column) -> TypedColumn:
@@ -412,10 +410,7 @@ def map_unresolved_function(
     def _resolve_aggregate_exp(
         result_exp: Column, default_result_type: DataType
     ) -> TypedColumn:
-        if is_in_pivot():
-            # it's not possible to cast the result in pivot
-            return _type_with_typer(result_exp)
-        elif is_window_enabled():
+        if is_window_enabled():
             # defer casting to capture whole window expression
             return TypedColumnWithDeferredCast(
                 result_exp, lambda: [default_result_type]
@@ -518,7 +513,6 @@ def map_unresolved_function(
         STRUCTURED_INCOMPATIBLE_AGGREGATES = {"min", "max"}
         if (
             aggregate_func.__name__ in STRUCTURED_INCOMPATIBLE_AGGREGATES
-            and not is_in_pivot()
             and not is_window_enabled()
             and isinstance(typed_arg.typ, (ArrayType, MapType, StructType))
         ):
@@ -526,12 +520,7 @@ def map_unresolved_function(
             variant_arg = snowpark_fn.to_variant(typed_arg.col)
             result = aggregate_func(variant_arg)
 
-            if is_in_pivot():
-                # In PIVOT context, casting is not allowed, so return VARIANT type
-                return _type_with_typer(result)
-            else:
-                # Cast back to the original type to maintain the expected schema
-                return TypedColumn(result.cast(typed_arg.typ), lambda: expected_types)
+            return TypedColumn(result.cast(typed_arg.typ), lambda: expected_types)
         else:
             # No structured type conversion needed
             result = aggregate_func(typed_arg.col)
@@ -912,15 +901,28 @@ def map_unresolved_function(
                 ):
                     # String + YearMonthInterval: Spark tries to cast string to double first, throws error if it fails
                     result_type = StringType()
+                    raise_error = _raise_error_helper(StringType(), AnalysisException)
                     if isinstance(snowpark_typed_args[0].typ, StringType):
-                        result_exp = (
-                            snowpark_fn.cast(snowpark_args[0], "double")
-                            + snowpark_args[1]
-                        )
+                        # Try to cast string to double, if it fails (returns null), raise exception
+                        cast_result = snowpark_fn.try_cast(snowpark_args[0], "double")
+                        result_exp = snowpark_fn.when(
+                            cast_result.is_null(),
+                            raise_error(
+                                snowpark_fn.lit(
+                                    f'The value \'{snowpark_args[0]}\' of the type {snowpark_typed_args[0].typ} cannot be cast to "DOUBLE" because it is malformed. Correct the value as per the syntax, or change its target type. Use `try_cast` to tolerate malformed input and return NULL instead. If necessary set "spark.sql.ansi.enabled" to "false" to bypass this error.'
+                                )
+                            ),
+                        ).otherwise(cast_result + snowpark_args[1])
                     else:
-                        result_exp = snowpark_args[0] + snowpark_fn.cast(
-                            snowpark_args[1], "double"
-                        )
+                        cast_result = snowpark_fn.try_cast(snowpark_args[1], "double")
+                        result_exp = snowpark_fn.when(
+                            cast_result.is_null(),
+                            raise_error(
+                                snowpark_fn.lit(
+                                    f'The value \'{snowpark_args[0]}\' of the type {snowpark_typed_args[0].typ} cannot be cast to "DOUBLE" because it is malformed. Correct the value as per the syntax, or change its target type. Use `try_cast` to tolerate malformed input and return NULL instead. If necessary set "spark.sql.ansi.enabled" to "false" to bypass this error.'
+                                )
+                            ),
+                        ).otherwise(snowpark_args[0] + cast_result)
                 case (StringType(), t) | (t, StringType()) if isinstance(
                     t, DayTimeIntervalType
                 ):
@@ -1016,9 +1018,12 @@ def map_unresolved_function(
                     result_exp = snowpark_args[0] - snowpark_args[1]
                     result_exp = result_exp.cast(result_type)
                 case (DateType(), DateType()):
-                    # TODO SNOW-2034420: resolve return type (it should be INTERVAL DAY)
-                    result_type = LongType()
-                    result_exp = snowpark_args[0] - snowpark_args[1]
+                    result_type = DayTimeIntervalType(
+                        DayTimeIntervalType.DAY, DayTimeIntervalType.DAY
+                    )
+                    result_exp = snowpark_fn.interval_day_time_from_parts(
+                        snowpark_args[0] - snowpark_args[1]
+                    )
                 case (DateType(), DayTimeIntervalType()) | (
                     DateType(),
                     YearMonthIntervalType(),
@@ -1037,14 +1042,22 @@ def map_unresolved_function(
                         result_type = TimestampType()
                         result_exp = snowpark_args[0] - snowpark_args[1]
                     else:
-                        # TODO SNOW-2034420: resolve return type (it should be INTERVAL DAY)
-                        result_type = LongType()
                         input_type = (
                             DateType() if spark_sql_ansi_enabled else DoubleType()
                         )
-                        result_exp = snowpark_args[0] - snowpark_args[1].cast(
-                            input_type
-                        )
+                        if isinstance(input_type, DateType):
+                            result_type = DayTimeIntervalType(
+                                DayTimeIntervalType.DAY, DayTimeIntervalType.DAY
+                            )
+                            result_exp = snowpark_fn.interval_day_time_from_parts(
+                                snowpark_args[0] - snowpark_args[1].cast(input_type)
+                            )
+                        else:
+                            # If ANSI is disabled, cast to DoubleType and return long (legacy behavior)
+                            result_type = LongType()
+                            result_exp = snowpark_args[0] - snowpark_args[1].cast(
+                                input_type
+                            )
                 case (TimestampType(), DayTimeIntervalType()) | (
                     TimestampType(),
                     YearMonthIntervalType(),
@@ -1063,10 +1076,12 @@ def map_unresolved_function(
                             f'[DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE] Cannot resolve "{spark_function_name}" due to data type mismatch: Parameter 2 requires the ("INTERVAL") type for timestamp operations, however "{snowpark_arg_names[1]}" has the type "{snowpark_typed_args[1].typ}".',
                         )
                 case (StringType(), DateType()):
-                    # TODO SNOW-2034420: resolve return type (it should be INTERVAL DAY)
-                    result_type = LongType()
-                    input_type = DateType()
-                    result_exp = snowpark_args[0].cast(input_type) - snowpark_args[1]
+                    result_type = DayTimeIntervalType(
+                        DayTimeIntervalType.DAY, DayTimeIntervalType.DAY
+                    )
+                    result_exp = snowpark_fn.interval_day_time_from_parts(
+                        snowpark_args[0].cast(DateType()) - snowpark_args[1]
+                    )
                 case (DateType(), (IntegerType() | ShortType() | ByteType())):
                     result_type = DateType()
                     result_exp = snowpark_args[0] - snowpark_args[1]
@@ -6184,15 +6199,19 @@ def map_unresolved_function(
                 or isinstance(arg_type, DayTimeIntervalType)
                 else DoubleType()
             )
-
-        case "pow":
-            result_exp = snowpark_fn.pow(snowpark_args[0], snowpark_args[1])
-            result_type = DoubleType()
-        case "power":
-            spark_function_name = (
-                f"POWER({snowpark_arg_names[0]}, {snowpark_arg_names[1]})"
-            )
-            result_exp = snowpark_fn.pow(snowpark_args[0], snowpark_args[1])
+        case "pow" | "power":
+            spark_function_name = f"{function_name if function_name == 'pow' else function_name.upper()}({snowpark_arg_names[0]}, {snowpark_arg_names[1]})"
+            if not spark_sql_ansi_enabled:
+                snowpark_args = _validate_numeric_args(
+                    function_name, snowpark_typed_args, snowpark_args
+                )
+            result_exp = snowpark_fn.when(
+                snowpark_fn.equal_nan(snowpark_fn.cast(snowpark_args[0], FloatType()))
+                | snowpark_fn.equal_nan(
+                    snowpark_fn.cast(snowpark_args[1], FloatType())
+                ),
+                NAN,
+            ).otherwise(snowpark_fn.pow(snowpark_args[0], snowpark_args[1]))
             result_type = DoubleType()
         case "product":
             col = snowpark_args[0]
@@ -6939,10 +6958,10 @@ def map_unresolved_function(
             )
             result_type = ArrayType(ArrayType(StringType()))
         case "sequence":
-            if snowpark_typed_args[0].typ != snowpark_typed_args[1].typ or (
-                not isinstance(snowpark_typed_args[0].typ, _IntegralType)
-                or not isinstance(snowpark_typed_args[1].typ, _IntegralType)
-            ):
+            both_integral = isinstance(
+                snowpark_typed_args[0].typ, _IntegralType
+            ) and isinstance(snowpark_typed_args[1].typ, _IntegralType)
+            if not both_integral:
                 exception = AnalysisException(
                     f"""[DATATYPE_MISMATCH.SEQUENCE_WRONG_INPUT_TYPES] Cannot resolve "sequence({snowpark_arg_names[0]}, {snowpark_arg_names[1]})" due to data type mismatch: `sequence` uses the wrong parameter type. The parameter type must conform to:
                         1. The start and stop expressions must resolve to the same type.
@@ -9458,15 +9477,21 @@ def map_unresolved_function(
                 result_exp = snowpark_fn.year(snowpark_fn.to_date(snowpark_args[0]))
             result_type = LongType()
         case binary_method if binary_method in ("to_binary", "try_to_binary"):
-            binary_format = "hex"
+            binary_format = snowpark_fn.lit("hex")
+            arg_str = snowpark_fn.cast(snowpark_args[0], StringType())
             if len(snowpark_args) > 1:
                 binary_format = snowpark_args[1]
             result_exp = snowpark_fn.when(
                 snowpark_args[0].isNull(), snowpark_fn.lit(None)
             ).otherwise(
                 snowpark_fn.function(binary_method)(
-                    snowpark_fn.cast(snowpark_args[0], StringType()), binary_format
-                ),
+                    snowpark_fn.when(
+                        (snowpark_fn.length(arg_str) % 2 == 1)
+                        & (snowpark_fn.lower(binary_format) == snowpark_fn.lit("hex")),
+                        snowpark_fn.concat(snowpark_fn.lit("0"), arg_str),
+                    ).otherwise(arg_str),
+                    binary_format,
+                )
             )
             result_type = BinaryType()
         case udtf_name if udtf_name.lower() in session._udtfs:
@@ -10705,12 +10730,18 @@ def _try_sum_helper(
                     return snowpark_fn.lit(None), new_type
                 else:
                     non_null_rows = snowpark_fn.count(col_name)
-                    return aggregate_sum / non_null_rows, new_type
+                    # Use _divnull to handle case when non_null_rows is 0
+                    return _divnull(aggregate_sum, non_null_rows), new_type
             else:
                 new_type = DecimalType(
                     precision=min(38, arg_type.precision + 10), scale=arg_type.scale
                 )
-                return aggregate_sum, new_type
+                # Return NULL when there are no non-null values (i.e., all values are NULL); this is handled using case/when to check for non-null values for both SUM and the sum component of AVG calculations.
+                non_null_rows = snowpark_fn.count(col_name)
+                result = snowpark_fn.when(
+                    non_null_rows == 0, snowpark_fn.lit(None)
+                ).otherwise(aggregate_sum)
+                return result, new_type
 
         case _:
             # If the input column is floating point (double and float are synonymous in Snowflake per
@@ -10728,9 +10759,16 @@ def _try_sum_helper(
                     return snowpark_fn.lit(None), DoubleType()
                 else:
                     non_null_rows = snowpark_fn.count(col_name)
-                    return aggregate_sum / non_null_rows, DoubleType()
+                    # Use _divnull to handle case when non_null_rows is 0
+                    return _divnull(aggregate_sum, non_null_rows), DoubleType()
             else:
-                return aggregate_sum, DoubleType()
+                # When all values are NULL, SUM should return NULL (not 0)
+                # Use case/when to return NULL when there are no non-null values (i.e., all values are NULL)
+                non_null_rows = snowpark_fn.count(col_name)
+                result = snowpark_fn.when(
+                    non_null_rows == 0, snowpark_fn.lit(None)
+                ).otherwise(aggregate_sum)
+                return result, DoubleType()
 
 
 def _get_type_precision(typ: DataType) -> tuple[int, int]: