snowpark-connect 0.26.0__py3-none-any.whl → 0.28.0__py3-none-any.whl

snowflake/snowpark_connect/column_name_handler.py

@@ -20,7 +20,6 @@ from snowflake.snowpark._internal.analyzer.analyzer_utils import (
 from snowflake.snowpark._internal.utils import quote_name
 from snowflake.snowpark.types import StructType
 from snowflake.snowpark_connect.config import global_config
-from snowflake.snowpark_connect.hidden_column import HiddenColumn
 from snowflake.snowpark_connect.utils.context import get_current_operation_scope
 from snowflake.snowpark_connect.utils.identifiers import (
     split_fully_qualified_spark_name,
@@ -124,7 +123,6 @@ class ColumnNameMap:
         ] = lambda: global_config.spark_sql_caseSensitive,
         column_metadata: dict | None = None,
         column_qualifiers: list[list[str]] | None = None,
-        hidden_columns: set[HiddenColumn] | None = None,
         parent_column_name_map: ColumnNameMap | None = None,
     ) -> None:
         """
@@ -135,7 +133,6 @@ class ColumnNameMap:
         The key is the original Spark column name, and the value is the metadata.
         example: Dict('age', {'foo': 'bar'})
         column_qualifiers: Optional qualifiers for the columns, used to handle table aliases or DataFrame aliases.
-        hidden_columns: Optional set of HiddenColumn objects.
         parent_column_name_map: parent ColumnNameMap
         """
         self.columns: list[ColumnNames] = []
@@ -144,7 +141,6 @@ class ColumnNameMap:
         self.snowpark_to_col = defaultdict(list)
         self.is_case_sensitive = is_case_sensitive
         self.column_metadata = column_metadata
-        self.hidden_columns = hidden_columns
 
         # Rename chain dictionary to track column renaming history
         self.rename_chains: dict[str, str] = {}  # old_name -> new_name mapping
@@ -338,8 +334,6 @@ class ColumnNameMap:
         *,
         allow_non_exists: bool = False,
         return_first: bool = False,
-        is_qualified: bool = False,
-        source_qualifiers: list[str] | None = None,
     ) -> str | None:
         assert isinstance(spark_column_name, str)
         resolved_name = (
@@ -347,37 +341,9 @@ class ColumnNameMap:
             if self.rename_chains
             else spark_column_name
         )
-
-        # We need to check hidden columns first. We want to avoid the code path
-        # within get_snowpark_column_names_from_spark_column_names that checks the parent ColumnNameMap.
-        # This is because that will return the name of the using column that's been dropped from the result
-        # dataframe. We want to fetch and resolve the hidden column to its visible using column name instead.
-        # Even if this is an unqualified reference or one to the visible column, it will resolve correctly to
-        # the visible name anyway.
-        snowpark_names = []
-        # Only check hidden columns for qualified references with source qualifiers
-        if is_qualified and source_qualifiers is not None and self.hidden_columns:
-            column_name = spark_column_name
-
-            # Check each hidden column for column name AND qualifier match
-            for hidden_col in self.hidden_columns:
-                if (
-                    hidden_col.spark_name == column_name
-                    and hidden_col.qualifiers == source_qualifiers
-                ):
-                    if not global_config.spark_sql_caseSensitive:
-                        if hidden_col.spark_name.upper() == column_name.upper() and [
-                            q.upper() for q in hidden_col.qualifiers
-                        ] == [q.upper() for q in source_qualifiers]:
-                            snowpark_names.append(hidden_col.visible_snowpark_name)
-                    else:
-                        snowpark_names.append(hidden_col.visible_snowpark_name)
-
-        # If not found in hidden columns, proceed with normal lookup
-        if not snowpark_names:
-            snowpark_names = self.get_snowpark_column_names_from_spark_column_names(
-                [resolved_name], return_first
-            )
+        snowpark_names = self.get_snowpark_column_names_from_spark_column_names(
+            [resolved_name], return_first
+        )
 
         snowpark_names_len = len(snowpark_names)
         if snowpark_names_len > 1:
@@ -464,27 +430,6 @@ class ColumnNameMap:
                 snowpark_columns.append(c.snowpark_name)
                 qualifiers.append(c.qualifiers)
 
-        # Note: The following code is commented out because there is a bug with handling duplicate columns in
-        # qualified select *'s. This needs to be revisited once a solution for that is found.
-        # TODO: https://snowflakecomputing.atlassian.net/browse/SNOW-2265240
-
-        # # Handles fetching/resolving the hidden columns if they also match the qualifiers
-        # # This method is only ever called for qualified references, so we need to check hidden columns as well.
-        # if self.hidden_columns:
-        #     for hidden_col in self.hidden_columns:
-        #         col_qualifiers = (
-        #             [q.upper() for q in hidden_col.qualifiers]
-        #             if not self.is_case_sensitive()
-        #             else hidden_col.qualifiers
-        #         )
-        #         if len(col_qualifiers) < len(qualifiers_input):
-        #             continue
-        #         if col_qualifiers[-len(qualifiers_input) :] == qualifiers_input:
-        #             # This hidden column matches! Add it to the results
-        #             spark_columns.append(hidden_col.spark_name)
-        #             snowpark_columns.append(hidden_col.visible_snowpark_name)
-        #             qualifiers.append(hidden_col.qualifiers)
-
         return spark_columns, snowpark_columns, qualifiers
 
     def get_snowpark_columns(self) -> list[str]:
@@ -616,35 +561,6 @@ class ColumnNameMap:
         else:
             return spark_name.upper()
 
-    def is_hidden_column_reference(
-        self, spark_column_name: str, source_qualifiers: list[str] | None = None
-    ) -> bool:
-        """
-        Check if a column reference would be resolved through hidden columns.
-        """
-        if not self.hidden_columns or source_qualifiers is None:
-            return False
-
-        # For qualified references with source_qualifiers
-        column_name = (
-            spark_column_name  # When has_plan_id=True, this is just the column name
-        )
-
-        for hidden_col in self.hidden_columns:
-            if (
-                hidden_col.spark_name == column_name
-                and hidden_col.qualifiers == source_qualifiers
-            ):
-                if not global_config.spark_sql_caseSensitive:
-                    if hidden_col.spark_name.upper() == column_name.upper() and [
-                        q.upper() for q in hidden_col.qualifiers
-                    ] == [q.upper() for q in source_qualifiers]:
-                        return True
-                else:
-                    return True
-
-        return False
-
 
 class JoinColumnNameMap(ColumnNameMap):
     def __init__(
@@ -654,9 +570,6 @@ class JoinColumnNameMap(ColumnNameMap):
     ) -> None:
         self.left_column_mapping: ColumnNameMap = left_colmap
         self.right_column_mapping: ColumnNameMap = right_colmap
-        # Ensure attributes expected by base-class helpers exist to avoid AttributeError
-        # when generic code paths (e.g., hidden column checks) touch them.
-        self.hidden_columns: set[HiddenColumn] | None = None
 
     def get_snowpark_column_name_from_spark_column_name(
         self,
@@ -664,9 +577,6 @@ class JoinColumnNameMap(ColumnNameMap):
         *,
         allow_non_exists: bool = False,
         return_first: bool = False,
-        # JoinColumnNameMap will never be called with using columns, so these parameters are not used.
-        is_qualified: bool = False,
-        source_qualifiers: list[str] | None = None,
     ) -> str | None:
         snowpark_column_name_in_left = (
             self.left_column_mapping.get_snowpark_column_name_from_spark_column_name(

snowflake/snowpark_connect/config.py

@@ -8,7 +8,7 @@ import re
 import sys
 from collections import defaultdict
 from copy import copy, deepcopy
-from typing import Any
+from typing import Any, Dict
 
 import jpype
 import pyspark.sql.connect.proto.base_pb2 as proto_base
@@ -17,6 +17,7 @@ from tzlocal import get_localzone_name
 from snowflake import snowpark
 from snowflake.snowpark._internal.analyzer.analyzer_utils import (
     quote_name_without_upper_casing,
+    unquote_if_quoted,
 )
 from snowflake.snowpark.exceptions import SnowparkSQLException
 from snowflake.snowpark.types import TimestampTimeZone, TimestampType
@@ -171,9 +172,6 @@ class GlobalConfig:
         "spark.app.name": lambda session, name: setattr(
             session, "query_tag", f"Spark-Connect-App-Name={name}"
         ),
-        "snowpark.connect.udf.packages": lambda session, packages: session.add_packages(
-            *packages.strip("[] ").split(",")
-        ),
         "snowpark.connect.udf.imports": lambda session, imports: parse_imports(
             session, imports
         ),
@@ -260,6 +258,7 @@ SESSION_CONFIG_KEY_WHITELIST = {
     "spark.sql.execution.pythonUDTF.arrow.enabled",
     "spark.sql.tvf.allowMultipleTableArguments.enabled",
     "snowpark.connect.sql.passthrough",
+    "snowpark.connect.cte.optimization_enabled",
     "snowpark.connect.iceberg.external_volume",
     "snowpark.connect.sql.identifiers.auto-uppercase",
     "snowpark.connect.udtf.compatibility_mode",
@@ -284,6 +283,7 @@ class SessionConfig:
     default_session_config = {
         "snowpark.connect.sql.identifiers.auto-uppercase": "all_except_columns",
         "snowpark.connect.sql.passthrough": "false",
+        "snowpark.connect.cte.optimization_enabled": "true",
         "snowpark.connect.udtf.compatibility_mode": "false",
         "snowpark.connect.views.duplicate_column_names_handling_mode": "rename",
         "spark.sql.execution.pythonUDTF.arrow.enabled": "false",
@@ -293,6 +293,7 @@ class SessionConfig:
 
     def __init__(self) -> None:
         self.config = deepcopy(self.default_session_config)
+        self.table_metadata: Dict[str, Dict[str, Any]] = {}
 
     def __getitem__(self, item: str) -> str:
         return self.get(item)
@@ -572,6 +573,12 @@ def set_snowflake_parameters(
                         snowpark_session.use_database(db)
                     case (prev, curr) if prev != curr:
                         snowpark_session.use_schema(prev)
+        case "snowpark.connect.cte.optimization_enabled":
+            # Set CTE optimization on the snowpark session
+            cte_enabled = str_to_bool(value)
+            snowpark_session.cte_optimization_enabled = cte_enabled
+            logger.info(f"Updated snowpark session CTE optimization: {cte_enabled}")
+
         case _:
             pass
 
@@ -581,6 +588,16 @@ def get_boolean_session_config_param(name: str) -> bool:
     return str_to_bool(session_config[name])
 
 
+def get_string_session_config_param(name: str) -> str:
+    session_config = sessions_config[get_session_id()]
+    return str(session_config[name])
+
+
+def get_cte_optimization_enabled() -> bool:
+    """Get the CTE optimization configuration setting."""
+    return get_boolean_session_config_param("snowpark.connect.cte.optimization_enabled")
+
+
 def auto_uppercase_column_identifiers() -> bool:
     session_config = sessions_config[get_session_id()]
     return session_config[
@@ -616,3 +633,81 @@ def get_timestamp_type():
             # shouldn't happen since `spark.sql.timestampType` is always defined, and `spark.conf.unset` sets it to default (TIMESTAMP_LTZ)
             timestamp_type = TimestampType(TimestampTimeZone.LTZ)
     return timestamp_type
+
+
+def record_table_metadata(
+    table_identifier: str,
+    table_type: str,
+    data_source: str,
+    supports_column_rename: bool = True,
+) -> None:
+    """
+    Record metadata about a table for Spark compatibility checks.
+
+    Args:
+        table_identifier: Full table identifier (catalog.database.table)
+        table_type: "v1" or "v2"
+        data_source: Source format (parquet, csv, iceberg, etc.)
+        supports_column_rename: Whether the table supports RENAME COLUMN
+    """
+    session_id = get_session_id()
+    session_config = sessions_config[session_id]
+
+    # Normalize table identifier for consistent lookup
+    # Use the full catalog.database.table identifier to avoid conflicts
+    normalized_identifier = table_identifier.upper().strip('"')
+
+    session_config.table_metadata[normalized_identifier] = {
+        "table_type": table_type,
+        "data_source": data_source,
+        "supports_column_rename": supports_column_rename,
+    }
+
+
+def get_table_metadata(table_identifier: str) -> Dict[str, Any] | None:
+    """
+    Get stored metadata for a table.
+
+    Args:
+        table_identifier: Full table identifier (catalog.database.table)
+
+    Returns:
+        Table metadata dict or None if not found
+    """
+    session_id = get_session_id()
+    session_config = sessions_config[session_id]
+
+    normalized_identifier = unquote_if_quoted(table_identifier).upper()
+
+    return session_config.table_metadata.get(normalized_identifier)
+
+
+def check_table_supports_operation(table_identifier: str, operation: str) -> bool:
+    """
+    Check if a table supports a given operation based on metadata and config.
+
+    Args:
+        table_identifier: Full table identifier (catalog.database.table)
+        operation: Operation to check (e.g., "rename_column")
+
+    Returns:
+        True if operation is supported, False if should be blocked
+    """
+    table_metadata = get_table_metadata(table_identifier)
+
+    if not table_metadata:
+        return True
+
+    session_id = get_session_id()
+    session_config = sessions_config[session_id]
+    enable_extensions = str_to_bool(
+        session_config.get("enable_snowflake_extension_behavior", "false")
+    )
+
+    if enable_extensions:
+        return True
+
+    if operation == "rename_column":
+        return table_metadata.get("supports_column_rename", True)
+
+    return True

snowflake/snowpark_connect/dataframe_container.py

@@ -8,7 +8,6 @@ from typing import TYPE_CHECKING, Callable
 
 from snowflake import snowpark
 from snowflake.snowpark.types import StructField, StructType
-from snowflake.snowpark_connect.hidden_column import HiddenColumn
 
 if TYPE_CHECKING:
     from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
@@ -61,7 +60,6 @@ class DataFrameContainer:
         column_metadata: dict | None = None,
         column_qualifiers: list[list[str]] | None = None,
         parent_column_name_map: ColumnNameMap | None = None,
-        hidden_columns: set[HiddenColumn] | None = None,
         table_name: str | None = None,
         alias: str | None = None,
         cached_schema_getter: Callable[[], StructType] | None = None,
@@ -78,7 +76,6 @@ class DataFrameContainer:
             column_metadata: Optional metadata dictionary
             column_qualifiers: Optional column qualifiers
             parent_column_name_map: Optional parent column name map
-            hidden_columns: Optional list of hidden column names
             table_name: Optional table name
             alias: Optional alias
             cached_schema_getter: Optional function to get cached schema
@@ -101,7 +98,6 @@ class DataFrameContainer:
             column_metadata,
             column_qualifiers,
             parent_column_name_map,
-            hidden_columns,
         )
 
         # Determine the schema getter to use
@@ -226,7 +222,6 @@ class DataFrameContainer:
         column_metadata: dict | None = None,
         column_qualifiers: list[list[str]] | None = None,
         parent_column_name_map: ColumnNameMap | None = None,
-        hidden_columns: set[HiddenColumn] | None = None,
     ) -> ColumnNameMap:
         """Create a ColumnNameMap with the provided configuration."""
         from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
@@ -237,7 +232,6 @@ class DataFrameContainer:
             column_metadata=column_metadata,
             column_qualifiers=column_qualifiers,
             parent_column_name_map=parent_column_name_map,
-            hidden_columns=hidden_columns,
         )
 
     @staticmethod

snowflake/snowpark_connect/expression/map_expression.py

@@ -6,6 +6,7 @@ import datetime
 from collections import defaultdict
 
 import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
+from pyspark.errors.exceptions.connect import AnalysisException
 
 import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
@@ -34,8 +35,10 @@ from snowflake.snowpark_connect.type_mapping import (
 from snowflake.snowpark_connect.typed_column import TypedColumn
 from snowflake.snowpark_connect.utils.context import (
     gen_sql_plan_id,
+    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,
@@ -134,7 +137,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
@@ -271,6 +277,30 @@ def map_expression(
         case "unresolved_function":
             return map_func.map_unresolved_function(exp, column_mapping, typer)
         case "unresolved_named_lambda_variable":
+            # Validate that this lambda variable is in scope
+            var_name = exp.unresolved_named_lambda_variable.name_parts[0]
+            current_params = get_current_lambda_params()
+
+            if current_params and var_name not in current_params:
+                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:
+                    raise 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."
+                    )
+
             col = snowpark_fn.Column(
                 UnresolvedAttribute(exp.unresolved_named_lambda_variable.name_parts[0])
             )

snowflake/snowpark_connect/expression/map_sql_expression.py

@@ -11,9 +11,10 @@ import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 import pyspark.sql.connect.proto.types_pb2 as types_proto
 from google.protobuf.any_pb2 import Any
 from pyspark.errors.exceptions.base import AnalysisException
+from pyspark.sql.connect import functions as pyspark_functions
 
 import snowflake.snowpark_connect.proto.snowflake_expression_ext_pb2 as snowflake_proto
-from snowflake import snowpark
+from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
 from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.typed_column import TypedColumn
@@ -89,6 +90,11 @@ def as_scala_seq(input):
     )
 
 
+@cache
+def _scala_some():
+    return jpype.JClass("scala.Some")
+
+
 def map_sql_expr(
     exp: expressions_proto.Expression,
     column_mapping: ColumnNameMap,
@@ -223,9 +229,6 @@ def apply_filter_clause(
 
 
 def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Expression:
-    from snowflake.snowpark_connect.expression.map_expression import (
-        map_single_column_expression,
-    )
     from snowflake.snowpark_connect.relation.map_sql import map_logical_plan_relation
 
     class_name = str(exp.getClass().getSimpleName())
@@ -308,22 +311,23 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
             )
             proto = expressions_proto.Expression(extension=any_proto)
         case "ExpressionWithUnresolvedIdentifier":
-            plan_id = None
-            identifierExpr = map_logical_plan_expression(exp.identifierExpr())
-            session = snowpark.Session.get_active_session()
-            m = ColumnNameMap([], [], None)
-            expr = map_single_column_expression(
-                identifierExpr, m, ExpressionTyper.dummy_typer(session)
+            from snowflake.snowpark_connect.relation.map_sql import (
+                get_relation_identifier_name,
             )
-            value = session.range(1).select(expr[1].col).collect()[0][0]
 
-            proto = expressions_proto.Expression(
-                unresolved_attribute=expressions_proto.Expression.UnresolvedAttribute(
-                    unparsed_identifier=str(value),
-                    plan_id=plan_id,
-                ),
-            )
-            # TODO: support identifier referencing unresolved function
+            value = unquote_if_quoted(get_relation_identifier_name(exp))
+            if getattr(pyspark_functions, value.lower(), None) is not None:
+                unresolved_function = exp.exprBuilder().apply(
+                    _scala_some()(value).toList()
+                )
+                proto = map_logical_plan_expression(unresolved_function)
+            else:
+                proto = expressions_proto.Expression(
+                    unresolved_attribute=expressions_proto.Expression.UnresolvedAttribute(
+                        unparsed_identifier=str(value),
+                        plan_id=None,
+                    ),
+                )
         case "InSubquery":
             rel_proto = map_logical_plan_relation(exp.query().plan())
             any_proto = Any()

snowflake/snowpark_connect/expression/map_unresolved_attribute.py

@@ -22,6 +22,7 @@ from snowflake.snowpark_connect.utils.context import (
     get_is_evaluating_sql,
     get_outer_dataframes,
     get_plan_id_map,
+    is_lambda_being_resolved,
     resolve_lca_alias,
 )
 from snowflake.snowpark_connect.utils.identifiers import (
@@ -162,7 +163,6 @@ def map_unresolved_attribute(
     attr_name = ".".join(name_parts)
 
     has_plan_id = exp.unresolved_attribute.HasField("plan_id")
-    source_qualifiers = None
 
     if has_plan_id:
         plan_id = exp.unresolved_attribute.plan_id
@@ -171,27 +171,13 @@ def map_unresolved_attribute(
         assert (
             target_df is not None
         ), f"resolving an attribute of a unresolved dataframe {plan_id}"
-
-        # Get the qualifiers for this column from the target DataFrame
-        source_qualifiers = (
-            target_df_container.column_map.get_qualifier_for_spark_column(
-                name_parts[-1]
-            )
-        )
-
-        if hasattr(column_mapping, "hidden_columns"):
-            hidden = column_mapping.hidden_columns
-        else:
-            hidden = None
-
         column_mapping = target_df_container.column_map
-        column_mapping.hidden_columns = hidden
         typer = ExpressionTyper(target_df)
 
-    def get_col(snowpark_name, has_hidden=False):
+    def get_col(snowpark_name):
         return (
             snowpark_fn.col(snowpark_name)
-            if not has_plan_id or has_hidden
+            if not has_plan_id
             else target_df.col(snowpark_name)
         )
 
@@ -276,17 +262,10 @@ def map_unresolved_attribute(
         quoted_attr_name = name_parts[0]
 
     snowpark_name = column_mapping.get_snowpark_column_name_from_spark_column_name(
-        quoted_attr_name,
-        allow_non_exists=True,
-        is_qualified=has_plan_id,
-        source_qualifiers=source_qualifiers if has_plan_id else None,
+        quoted_attr_name, allow_non_exists=True
     )
-
     if snowpark_name is not None:
-        is_hidden = column_mapping.is_hidden_column_reference(
-            quoted_attr_name, source_qualifiers
-        )
-        col = get_col(snowpark_name, is_hidden)
+        col = get_col(snowpark_name)
         qualifiers = column_mapping.get_qualifier_for_spark_column(quoted_attr_name)
     else:
         # this means it has to be a struct column with a field name
@@ -356,6 +335,23 @@ def map_unresolved_attribute(
                     return (unqualified_name, typed_col)
 
         if snowpark_name is None:
+            # Check if we're inside a lambda and trying to reference an outer column
+            # This catches direct column references (not lambda variables)
+            if is_lambda_being_resolved() and column_mapping:
+                # Check if this column exists in the outer scope (not lambda params)
+                outer_col_name = (
+                    column_mapping.get_snowpark_column_name_from_spark_column_name(
+                        attr_name, allow_non_exists=True
+                    )
+                )
+                if outer_col_name:
+                    # This is an outer scope column being referenced inside a lambda
+                    raise AnalysisException(
+                        f"Reference to non-lambda variable '{attr_name}' within lambda function. "
+                        f"Lambda functions can only access their own parameters. "
+                        f"If '{attr_name}' is a table column, it must be passed as an explicit parameter to the enclosing function."
+                    )
+
             if has_plan_id:
                 raise AnalysisException(
                     f'[RESOLVED_REFERENCE_COLUMN_NOT_FOUND] The column "{attr_name}" does not exist in the target dataframe.'

snowflake/snowpark_connect/expression/map_unresolved_function.py

@@ -711,6 +711,9 @@ def map_unresolved_function(
                 "-",
             )
             match (snowpark_typed_args[0].typ, snowpark_typed_args[1].typ):
+                case (DateType(), NullType()) | (NullType(), DateType()):
+                    result_type = LongType()
+                    result_exp = snowpark_fn.lit(None).cast(result_type)
                 case (NullType(), _) | (_, NullType()):
                     result_type = _get_add_sub_result_type(
                         snowpark_typed_args[0].typ,
@@ -724,7 +727,10 @@ def map_unresolved_function(
                     result_type = LongType()
                     result_exp = snowpark_args[0] - snowpark_args[1]
                 case (DateType(), StringType()):
-                    if "INTERVAL" == snowpark_typed_args[1].col._expr1.pretty_name:
+                    if (
+                        hasattr(snowpark_typed_args[1].col._expr1, "pretty_name")
+                        and "INTERVAL" == snowpark_typed_args[1].col._expr1.pretty_name
+                    ):
                         result_type = TimestampType()
                         result_exp = snowpark_args[0] - snowpark_args[1]
                     else:
@@ -2421,7 +2427,7 @@ def map_unresolved_function(
                 "try_to_date",
                 snowpark_fn.cast(
                     truncated_date,
-                    TimestampType(snowpark.types.TimestampTimeZone.NTZ),
+                    TimestampType(),
                 ),
                 snowpark_args[1],
             )
@@ -2613,9 +2619,18 @@ def map_unresolved_function(
                     result_type = input_type.element_type
                     result_exp = fn(snowpark_args[0])
                 case _:
-                    spark_col_names = ["key", "value"]
-                    result_exp = fn(snowpark_args[0])
-                    result_type = [input_type.key_type, input_type.value_type]
+                    # Check if the type has map-like attributes before accessing them
+                    if hasattr(input_type, "key_type") and hasattr(
+                        input_type, "value_type"
+                    ):
+                        spark_col_names = ["key", "value"]
+                        result_exp = fn(snowpark_args[0])
+                        result_type = [input_type.key_type, input_type.value_type]
+                    else:
+                        # Throw proper error for types without key_type/value_type attributes
+                        raise AnalysisException(
+                            f'[DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE] Cannot resolve "{function_name}({snowpark_arg_names[0]})" due to data type mismatch: Parameter 1 requires the ("ARRAY" or "MAP") type, however "{snowpark_arg_names[0]}" has the type "{str(input_type)}".'
+                        )
         case "expm1":
             spark_function_name = f"EXPM1({snowpark_arg_names[0]})"
             result_exp = snowpark_fn.exp(*snowpark_args) - 1
@@ -8725,7 +8740,7 @@ def _resolve_function_with_lambda(
         artificial_df = Session.get_active_session().create_dataframe([], schema)
         set_schema_getter(artificial_df, lambda: schema)
 
-        with resolving_lambda_function():
+        with resolving_lambda_function(names):
             return map_expression(
                 (
                     lambda_exp.lambda_function.function
@@ -9911,7 +9926,10 @@ def _get_spark_function_name(
             return f"({date_param_name1} {operation_op} {date_param_name2})"
         case (StringType(), DateType()):
             date_param_name2 = _get_literal_param_name(exp, 1, snowpark_arg_names[1])
-            if "INTERVAL" == col1.col._expr1.pretty_name:
+            if (
+                hasattr(col1.col._expr1, "pretty_name")
+                and "INTERVAL" == col1.col._expr1.pretty_name
+            ):
                 return f"{date_param_name2} {operation_op} {snowpark_arg_names[0]}"
             elif global_config.spark_sql_ansi_enabled and function_name == "+":
                 return f"{operation_func}(cast({date_param_name2} as date), cast({snowpark_arg_names[0]} as double))"
@@ -9919,9 +9937,9 @@ def _get_spark_function_name(
                 return f"({snowpark_arg_names[0]} {operation_op} {date_param_name2})"
         case (DateType(), StringType()):
             date_param_name1 = _get_literal_param_name(exp, 0, snowpark_arg_names[0])
-            if (
-                global_config.spark_sql_ansi_enabled
-                or "INTERVAL" == col2.col._expr1.pretty_name
+            if global_config.spark_sql_ansi_enabled or (
+                hasattr(col2.col._expr1, "pretty_name")
+                and "INTERVAL" == col2.col._expr1.pretty_name
             ):
                 return f"{date_param_name1} {operation_op} {snowpark_arg_names[1]}"
             else: