snowpark-connect 0.27.0__py3-none-any.whl → 1.7.0__py3-none-any.whl

snowflake/snowpark_connect/expression/map_cast.py

@@ -6,17 +6,22 @@ import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 import pyspark.sql.connect.proto.types_pb2 as types_proto
 from pyspark.errors.exceptions.base import (
     AnalysisException,
+    ArithmeticException,
+    IllegalArgumentException,
     NumberFormatException,
     SparkRuntimeException,
 )
 
 import snowflake.snowpark.functions as snowpark_fn
+from snowflake.snowpark.column import Column
 from snowflake.snowpark.types import (
     BinaryType,
     BooleanType,
     DataType,
     DateType,
+    DecimalType,
     DoubleType,
+    FloatType,
     IntegerType,
     LongType,
     MapType,
@@ -25,12 +30,22 @@ from snowflake.snowpark.types import (
     StructType,
     TimestampTimeZone,
     TimestampType,
+    YearMonthIntervalType,
     _FractionalType,
     _IntegralType,
     _NumericType,
 )
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
 from snowflake.snowpark_connect.config import global_config
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
+from snowflake.snowpark_connect.expression.error_utils import raise_error_helper
+from snowflake.snowpark_connect.expression.integral_types_support import (
+    apply_fractional_to_integral_cast,
+    apply_fractional_to_integral_cast_with_ansi_check,
+    apply_integral_overflow_with_ansi_check,
+    get_integral_type_bounds,
+)
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.type_mapping import (
     map_type_string_to_snowpark_type,
@@ -49,7 +64,7 @@ SYMBOL_FUNCTIONS = {"<", ">", "<=", ">=", "!=", "+", "-", "*", "/", "%", "div"}
 CAST_FUNCTIONS = {
     "boolean": types_proto.DataType(boolean=types_proto.DataType.Boolean()),
     "int": types_proto.DataType(integer=types_proto.DataType.Integer()),
-    "smallint": types_proto.DataType(integer=types_proto.DataType.Integer()),
+    "smallint": types_proto.DataType(short=types_proto.DataType.Short()),
     "bigint": types_proto.DataType(long=types_proto.DataType.Long()),
     "tinyint": types_proto.DataType(byte=types_proto.DataType.Byte()),
     "float": types_proto.DataType(float=types_proto.DataType.Float()),
@@ -87,7 +102,9 @@ def map_cast(
             to_type = map_type_string_to_snowpark_type(exp.cast.type_str)
             to_type_str = exp.cast.type_str.upper()
         case _:
-            raise ValueError("No type to cast to")
+            exception = ValueError("No type to cast to")
+            attach_custom_error_code(exception, ErrorCodes.INVALID_CAST)
+            raise exception
 
     from_exp = exp.cast.expr
     new_name, typed_column = map_single_column_expression(
@@ -230,6 +247,11 @@ def map_cast(
         case (_, BooleanType()) if isinstance(from_type, _NumericType):
             result_exp = col.cast(LongType()).cast(to_type)
 
+        case (_IntegralType(), _IntegralType()):
+            result_exp = apply_integral_overflow_with_ansi_check(
+                col, to_type, spark_sql_ansi_enabled
+            )
+
         # binary
         case (StringType(), BinaryType()):
             result_exp = snowpark_fn.to_binary(col, "UTF-8")
@@ -265,6 +287,44 @@ def map_cast(
             result_exp = snowpark_fn.to_varchar(col, "UTF-8")
 
         # numeric
+        case (_, _) if isinstance(from_type, (FloatType, DoubleType)) and isinstance(
+            to_type, _IntegralType
+        ):
+            truncated = (
+                snowpark_fn.when(
+                    col == snowpark_fn.lit(float("nan")), snowpark_fn.lit(0)
+                )
+                .when(col < 0, snowpark_fn.ceil(col))
+                .otherwise(snowpark_fn.floor(col))
+            )
+
+            if spark_sql_ansi_enabled:
+                result_exp = apply_fractional_to_integral_cast_with_ansi_check(
+                    truncated, to_type, True
+                )
+            else:
+                target_min, target_max = get_integral_type_bounds(to_type)
+                result_exp = (
+                    snowpark_fn.when(
+                        truncated > snowpark_fn.lit(target_max),
+                        snowpark_fn.lit(target_max),
+                    )
+                    .when(
+                        truncated < snowpark_fn.lit(target_min),
+                        snowpark_fn.lit(target_min),
+                    )
+                    .otherwise(truncated.cast(to_type))
+                )
+        case (_, _) if isinstance(from_type, DecimalType) and isinstance(
+            to_type, _IntegralType
+        ):
+            result_exp = snowpark_fn.when(col < 0, snowpark_fn.ceil(col)).otherwise(
+                snowpark_fn.floor(col)
+            )
+            result_exp = result_exp.cast(to_type)
+            result_exp = apply_integral_overflow_with_ansi_check(
+                result_exp, to_type, spark_sql_ansi_enabled
+            )
         case (_, _) if isinstance(from_type, _FractionalType) and isinstance(
             to_type, _IntegralType
         ):
@@ -275,16 +335,49 @@ def map_cast(
                 .when(col < 0, snowpark_fn.ceil(col))
                 .otherwise(snowpark_fn.floor(col))
             )
-            result_exp = result_exp.cast(to_type)
+            result_exp = apply_fractional_to_integral_cast(result_exp, to_type)
         case (StringType(), _) if (isinstance(to_type, _IntegralType)):
             if spark_sql_ansi_enabled:
-                result_exp = snowpark_fn.cast(col, DoubleType())
+                double_val = snowpark_fn.cast(col, DoubleType())
+
+                target_min, target_max = get_integral_type_bounds(to_type)
+                raise_error = raise_error_helper(to_type, NumberFormatException)
+                to_type_name = to_type.__class__.__name__.upper().replace("TYPE", "")
+
+                truncated = snowpark_fn.when(
+                    double_val < 0, snowpark_fn.ceil(double_val)
+                ).otherwise(snowpark_fn.floor(double_val))
+
+                result_exp = snowpark_fn.when(
+                    (truncated < snowpark_fn.lit(target_min))
+                    | (truncated > snowpark_fn.lit(target_max)),
+                    raise_error(
+                        snowpark_fn.lit("[CAST_INVALID_INPUT] The value '"),
+                        col,
+                        snowpark_fn.lit(
+                            f'\' of the type "STRING" cannot be cast to "{to_type_name}" 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(truncated.cast(to_type))
             else:
-                result_exp = snowpark_fn.try_cast(col, DoubleType())
-            result_exp = snowpark_fn.when(
-                result_exp < 0, snowpark_fn.ceil(result_exp)
-            ).otherwise(snowpark_fn.floor(result_exp))
-            result_exp = result_exp.cast(to_type)
+                double_val = snowpark_fn.try_cast(col, DoubleType())
+
+                truncated = snowpark_fn.when(
+                    double_val < 0, snowpark_fn.ceil(double_val)
+                ).otherwise(snowpark_fn.floor(double_val))
+
+                target_min, target_max = get_integral_type_bounds(to_type)
+                result_exp = (
+                    snowpark_fn.when(
+                        double_val.isNull(), snowpark_fn.lit(None).cast(to_type)
+                    )
+                    .when(
+                        (truncated < snowpark_fn.lit(target_min))
+                        | (truncated > snowpark_fn.lit(target_max)),
+                        snowpark_fn.lit(None).cast(to_type),
+                    )
+                    .otherwise(truncated.cast(to_type))
+                )
         # https://docs.snowflake.com/en/sql-reference/functions/try_cast Only works on certain types (mostly non-structured ones)
         case (StringType(), _) if isinstance(to_type, _NumericType) or isinstance(
             to_type, StringType
@@ -299,10 +392,16 @@ def map_cast(
                 result_exp = snowpark_fn.cast(col, to_type)
             else:
                 result_exp = snowpark_fn.try_cast(col, to_type)
+        case (StringType(), YearMonthIntervalType()):
+            result_exp = _cast_string_to_year_month_interval(col, to_type)
+        case (YearMonthIntervalType(), StringType()):
+            result_exp = _cast_year_month_interval_to_string(col, from_type)
         case (StringType(), _):
-            raise AnalysisException(
+            exception = AnalysisException(
                 f"""[DATATYPE_MISMATCH.CAST_WITHOUT_SUGGESTION] Cannot resolve "{col_name}" due to data type mismatch: cannot cast "{snowpark_to_proto_type(from_type, column_mapping)}" to "{exp.cast.type_str.upper()}".;"""
             )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_CAST)
+            raise exception
         case _:
             result_exp = snowpark_fn.cast(col, to_type)
 
@@ -317,9 +416,11 @@ def sanity_check(
     """
 
     if isinstance(from_type, LongType) and isinstance(to_type, BinaryType):
-        raise NumberFormatException(
+        exception = NumberFormatException(
             f"""[DATATYPE_MISMATCH.CAST_WITH_CONF_SUGGESTION] Cannot resolve "CAST({value} AS BINARY)" due to data type mismatch: cannot cast "BIGINT" to "BINARY" with ANSI mode on."""
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_CAST)
+        raise exception
 
     if (
         from_type_cast
@@ -329,9 +430,11 @@ def sanity_check(
         if value is not None:
             value = value.strip().lower()
         if value not in {"t", "true", "f", "false", "y", "yes", "n", "no", "0", "1"}:
-            raise SparkRuntimeException(
+            exception = SparkRuntimeException(
                 f"""[CAST_INVALID_INPUT] The value '{value}' of the type "STRING" cannot be cast to "BOOLEAN" 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."""
             )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_CAST)
+            raise exception
 
     raise_cast_failure_exception = False
     if isinstance(to_type, _IntegralType):
@@ -351,6 +454,112 @@ def sanity_check(
         except Exception:
             raise_cast_failure_exception = True
     if raise_cast_failure_exception:
-        raise NumberFormatException(
-            """[CAST_INVALID_INPUT] 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 setting "spark.sql.ansi.enabled" to "false" may bypass this error."""
-        )
+        if not isinstance(from_type, StringType) and isinstance(to_type, _IntegralType):
+            from_type_name = from_type.__class__.__name__.upper().replace("TYPE", "")
+            to_type_name = to_type.__class__.__name__.upper().replace("TYPE", "")
+            value_suffix = "L" if isinstance(from_type, LongType) else ""
+            exception = ArithmeticException(
+                f"""[CAST_OVERFLOW] The value {value}{value_suffix} of the type "{from_type_name}" cannot be cast to "{to_type_name}" due to an overflow. Use `try_cast` to tolerate overflow and return NULL instead. If necessary set "spark.sql.ansi.enabled" to "false" to bypass this error."""
+            )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_CAST)
+        else:
+            exception = NumberFormatException(
+                """[CAST_INVALID_INPUT] 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 setting "spark.sql.ansi.enabled" to "false" may bypass this error."""
+            )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_CAST)
+        raise exception
+
+
+def _cast_string_to_year_month_interval(col: Column, to_type: YearMonthIntervalType):
+    """
+    Cast string to year-month interval type.
+    Supports different interval formats:
+    1. '[+|-]y-m' format (e.g., '1-2', '-1-6', '+2-3') - YEAR TO MONTH
+    2. '[+|-]y' format (e.g., '1', '-2', '+3') - YEAR
+    3. '[+|-]m' format (e.g., '12', '-6', '+15') - MONTH
+    4. 'INTERVAL [+|-]'[+|-]y-m' YEAR TO MONTH' format - extract the y-m part
+    5. 'INTERVAL [+|-]'[+|-]y' YEAR' format - extract the y part
+    6. 'INTERVAL [+|-]'[+|-]m' MONTH' format - extract the m part
+    """
+    # Extract values from different formats
+    value = snowpark_fn.regexp_extract(col, "'([^']+)'", 1)
+    years = snowpark_fn.regexp_extract(col, "^[+-]?\\d+", 0)
+    months = snowpark_fn.regexp_extract(col, "-(\\d+)$", 1)
+    raise_error = raise_error_helper(to_type, IllegalArgumentException)
+
+    # For MONTH-only intervals, treat the input as months
+    if (
+        to_type.start_field == YearMonthIntervalType.MONTH
+        and to_type.end_field == YearMonthIntervalType.MONTH
+    ):
+        months = years
+        years = snowpark_fn.lit(0)
+
+    # Define overflow limits based on Snowflake's INTERVAL limits
+    # Maximum year-month interval is 178956970-7 (positive) and -178956970-8 (negative)
+    max_years = snowpark_fn.lit(178956970)
+    max_months_positive = snowpark_fn.lit(7)
+    max_months_negative = snowpark_fn.lit(8)
+
+    return snowpark_fn.when(
+        col.like("INTERVAL % YEAR TO MONTH")
+        | col.like("INTERVAL % YEAR")
+        | col.like("INTERVAL % MONTH"),
+        value.cast(to_type),
+    ).when(
+        col.rlike("^[+-]?\\d+(-\\d+)?$"),
+        snowpark_fn.when(
+            # Check for overflow conditions
+            ((years >= max_years) & (months > max_months_positive))
+            | (years > max_years)
+            | ((years <= -max_years) & (months > max_months_negative))
+            | (years < -max_years),
+            raise_error(snowpark_fn.lit("Error parsing interval year-month string")),
+        ).otherwise(col.cast(to_type)),
+    )
+
+
+def _cast_year_month_interval_to_string(col: Column, from_type: YearMonthIntervalType):
+    """
+    Cast year-month interval to string.
+    Returns format like 'INTERVAL '1-2' YEAR TO MONTH' for year-month intervals.
+    """
+    years = snowpark_fn.date_part("YEAR", col)
+    months = snowpark_fn.date_part("MONTH", col)
+
+    total_months = years * 12 + months
+
+    start_field = from_type.start_field  # YEAR
+    end_field = from_type.end_field  # MONTH
+
+    def _format_interval_udf(
+        total_months: int, start_field: int, end_field: int
+    ) -> str:
+        is_negative = total_months < 0
+        abs_months = abs(total_months)
+        years = abs_months // 12
+        months = abs_months % 12
+
+        is_year_only = start_field == 0 and end_field == 0
+        is_month_only = start_field == 1 and end_field == 1
+
+        if is_year_only:
+            sign = "-" if is_negative else ""
+            return f"INTERVAL '{sign}{years}' YEAR"
+        elif is_month_only:
+            return f"INTERVAL '{total_months}' MONTH"
+        else:  # YEAR TO MONTH
+            if is_negative:
+                return f"INTERVAL '-{years}-{months}' YEAR TO MONTH"
+            else:
+                return f"INTERVAL '{years}-{months}' YEAR TO MONTH"
+
+    format_udf = cached_udf(
+        _format_interval_udf,
+        input_types=[IntegerType(), IntegerType(), IntegerType()],
+        return_type=StringType(),
+    )
+
+    return format_udf(
+        total_months, snowpark_fn.lit(start_field), snowpark_fn.lit(end_field)
+    )

snowflake/snowpark_connect/expression/map_expression.py

@@ -14,6 +14,8 @@ from snowflake.snowpark import Session
 from snowflake.snowpark._internal.analyzer.expression import UnresolvedAttribute
 from snowflake.snowpark.types import TimestampTimeZone, TimestampType
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
 from snowflake.snowpark_connect.expression import (
     map_extension,
     map_udf,
@@ -38,6 +40,7 @@ from snowflake.snowpark_connect.utils.context import (
     get_current_lambda_params,
     is_function_argument_being_resolved,
     is_lambda_being_resolved,
+    not_resolving_fun_args,
 )
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
@@ -61,9 +64,11 @@ def map_alias(
         # Multi-column case: handle like explode("map").alias("key", "value")
         col_names, col = map_expression(alias.expr, column_mapping, typer)
         if len(col_names) != len(list(alias.name)):
-            raise ValueError(
+            exception = ValueError(
                 f"Found the unresolved operator: 'Project [{col_names} AS ({', '.join(list(alias.name))})]. Number of aliases ({len(list(alias.name))}) does not match number of columns ({len(col_names)})"
             )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+            raise exception
         return list(alias.name), col
 
     name, col = map_single_column_expression(alias.expr, column_mapping, typer)
@@ -136,7 +141,10 @@ def map_expression(
         case "expression_string":
             return map_sql_expr(exp, column_mapping, typer)
         case "extension":
-            return map_extension.map_extension(exp, column_mapping, typer)
+            # Extensions can be passed as function args, and we need to reset the context here.
+            # Matters only for resolving alias expressions in the extensions rel.
+            with not_resolving_fun_args():
+                return map_extension.map_extension(exp, column_mapping, typer)
         case "lambda_function":
             lambda_name, lambda_body = map_single_column_expression(
                 exp.lambda_function.function, column_mapping, typer
@@ -222,41 +230,39 @@ def map_expression(
                     | exp.sort_order.SORT_DIRECTION_ASCENDING
                 ):
                     if exp.sort_order.null_ordering == exp.sort_order.SORT_NULLS_LAST:
-                        return [child_name], snowpark_fn.asc_nulls_last(child_column)
+                        col = snowpark_fn.asc_nulls_last(child_column.col)
                     else:
                         # If nulls are not specified or null_ordering is FIRST in the sort order, Spark defaults to nulls
                         # first in the case of ascending sort order.
-                        return [child_name], snowpark_fn.asc_nulls_first(child_column)
+                        col = snowpark_fn.asc_nulls_first(child_column.col)
                 case exp.sort_order.SORT_DIRECTION_DESCENDING:
                     if exp.sort_order.null_ordering == exp.sort_order.SORT_NULLS_FIRST:
-                        return [child_name], snowpark_fn.desc_nulls_first(child_column)
+                        col = snowpark_fn.desc_nulls_first(child_column.col)
                     else:
                         # If nulls are not specified or null_ordering is LAST in the sort order, Spark defaults to nulls
                         # last in the case of descending sort order.
-                        return [child_name], snowpark_fn.desc_nulls_last(child_column)
+                        col = snowpark_fn.desc_nulls_last(child_column.col)
                 case _:
-                    raise ValueError(
+                    exception = ValueError(
                         f"Invalid sort direction {exp.sort_order.direction}"
                     )
+                    attach_custom_error_code(
+                        exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                    )
+                    raise exception
+            return [child_name], TypedColumn(col, lambda: typer.type(col))
         case "unresolved_attribute":
             col_name, col = map_att.map_unresolved_attribute(exp, column_mapping, typer)
             # Check if this is a multi-column regex expansion
             matched_cols = getattr(col, "_regex_matched_columns", list())
             if matched_cols:
                 # Create expressions for all matched columns
-                snowpark_cols = []
-                for matched_col in matched_cols:
-                    snowpark_name = (
-                        column_mapping.get_snowpark_column_name_from_spark_column_name(
-                            matched_col
-                        )
-                    )
-                    snowpark_cols.append(snowpark_name)
-
+                snowpark_cols = [c.snowpark_name for c in matched_cols]
+                spark_cols = [c.spark_name for c in matched_cols]
                 # Create a combined expression for all columns
                 col_expr = snowpark_fn.sql_expr(", ".join(snowpark_cols))
                 return (
-                    matched_cols,
+                    spark_cols,
                     TypedColumn(
                         col_expr,
                         lambda: [
@@ -271,6 +277,36 @@ def map_expression(
             )
             return [col_name], col
         case "unresolved_function":
+            from snowflake.snowpark_connect.utils.context import (
+                get_is_processing_order_by,
+            )
+
+            is_order_by = get_is_processing_order_by()
+            if is_order_by:
+                # For expressions in an order by clause check if we can reuse already-computed column.
+                if exp.unresolved_function.function_name:
+                    func_name = exp.unresolved_function.function_name
+                    available_columns = column_mapping.get_spark_columns()
+
+                    for col_name in available_columns:
+                        if (
+                            func_name.lower() in col_name.lower()
+                            and "(" in col_name
+                            and ")" in col_name
+                        ):
+                            # This looks like it might be an expression
+                            snowpark_col_name = column_mapping.get_snowpark_column_name_from_spark_column_name(
+                                col_name
+                            )
+                            if snowpark_col_name:
+                                # Optimization applied - reusing already computed column
+                                return [col_name], TypedColumn(
+                                    snowpark_fn.col(snowpark_col_name),
+                                    lambda col_name=snowpark_col_name: typer.type(
+                                        col_name
+                                    ),
+                                )
+
             return map_func.map_unresolved_function(exp, column_mapping, typer)
         case "unresolved_named_lambda_variable":
             # Validate that this lambda variable is in scope
@@ -278,13 +314,28 @@ def map_expression(
             current_params = get_current_lambda_params()
 
             if current_params and var_name not in current_params:
-                raise AnalysisException(
-                    f"Reference to non-lambda variable '{var_name}' within lambda function. "
-                    f"Lambda functions can only access their own parameters. "
-                    f"Available lambda parameters are: {current_params}. "
-                    f"If '{var_name}' is an outer scope lambda variable from a nested lambda, "
-                    f"that is an unsupported feature in Snowflake SQL."
+                outer_col_name = (
+                    column_mapping.get_snowpark_column_name_from_spark_column_name(
+                        var_name, allow_non_exists=True
+                    )
                 )
+                if outer_col_name:
+                    col = snowpark_fn.col(outer_col_name)
+                    return ["namedlambdavariable()"], TypedColumn(
+                        col, lambda: typer.type(col)
+                    )
+                else:
+                    exception = AnalysisException(
+                        f"Cannot resolve variable '{var_name}' within lambda function. "
+                        f"Lambda functions can access their own parameters and parent dataframe columns. "
+                        f"Current lambda parameters: {current_params}. "
+                        f"If '{var_name}' is an outer scope lambda variable from a nested lambda, "
+                        f"that is an unsupported feature in Snowflake SQL."
+                    )
+                    attach_custom_error_code(
+                        exception, ErrorCodes.UNSUPPORTED_OPERATION
+                    )
+                    raise exception
 
             col = snowpark_fn.Column(
                 UnresolvedAttribute(exp.unresolved_named_lambda_variable.name_parts[0])
@@ -296,10 +347,10 @@ def map_expression(
             p = exp.unresolved_regex.col_name
             pattern_str = p[1:-1] if p.startswith("`") and p.endswith("`") else p
 
-            mapping = column_mapping.spark_to_snowpark_for_pattern(pattern_str)
+            columns = column_mapping.get_columns_matching_pattern(pattern_str)
             spark_cols, snowpark_cols = (
-                [spark_name for spark_name, _ in mapping],
-                [snowpark_name for _, snowpark_name in mapping],
+                [c.spark_name for c in columns],
+                [c.snowpark_name for c in columns],
             )
 
             col_expr = snowpark_fn.sql_expr(", ".join(snowpark_cols))
@@ -319,6 +370,8 @@ def map_expression(
         case "update_fields":
             return map_update_fields.map_update_fields(exp, column_mapping, typer)
         case _:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported expression type {expr_type}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception