snowpark-connect 0.21.0__py3-none-any.whl → 0.23.0__py3-none-any.whl

snowflake/snowpark_connect/config.py

@@ -4,10 +4,8 @@
 
 # Proto source for reference:
 # https://github.com/apache/spark/blob/branch-3.5/connector/connect/common/src/main/protobuf/spark/connect/base.proto#L420
-import os
 import re
 import sys
-import time
 from collections import defaultdict
 from copy import copy, deepcopy
 from typing import Any
@@ -168,6 +166,9 @@ class GlobalConfig:
         "snowpark.connect.udf.packages": lambda session, packages: session.add_packages(
             *packages.strip("[] ").split(",")
         ),
+        "snowpark.connect.udf.imports": lambda session, imports: parse_imports(
+            session, imports
+        ),
     }
 
     float_config_list = []
@@ -332,7 +333,7 @@ def route_config_proto(
     match op_type:
         case "set":
             logger.info("SET")
-
+            telemetry.report_config_set(config.operation.set.pairs)
             for pair in config.operation.set.pairs:
                 # Check if the value field is present, not present when invalid fields are set in conf.
                 if not pair.HasField("value"):
@@ -342,7 +343,6 @@ def route_config_proto(
                         f"Cannot set config '{pair.key}' to None"
                     )
 
-                telemetry.report_config_set(pair.key, pair.value)
                 set_config_param(
                     config.session_id, pair.key, pair.value, snowpark_session
                 )
@@ -350,14 +350,15 @@ def route_config_proto(
             return proto_base.ConfigResponse(session_id=config.session_id)
         case "unset":
             logger.info("UNSET")
+            telemetry.report_config_unset(config.operation.unset.keys)
             for key in config.operation.unset.keys:
-                telemetry.report_config_unset(key)
                 unset_config_param(config.session_id, key, snowpark_session)
 
             return proto_base.ConfigResponse(session_id=config.session_id)
         case "get":
             logger.info("GET")
             res = proto_base.ConfigResponse(session_id=config.session_id)
+            telemetry.report_config_get(config.operation.get.keys)
             for key in config.operation.get.keys:
                 pair = res.pairs.add()
                 pair.key = key
@@ -367,6 +368,9 @@ def route_config_proto(
             return res
         case "get_with_default":
             logger.info("GET_WITH_DEFAULT")
+            telemetry.report_config_get(
+                [pair.key for pair in config.operation.get_with_default.pairs]
+            )
             result_pairs = [
                 proto_base.KeyValue(
                     key=pair.key,
@@ -383,6 +387,7 @@ def route_config_proto(
         case "get_option":
             logger.info("GET_OPTION")
             res = proto_base.ConfigResponse(session_id=config.session_id)
+            telemetry.report_config_get(config.operation.get_option.keys)
             for key in config.operation.get_option.keys:
                 pair = res.pairs.add()
                 pair.key = key
@@ -411,6 +416,7 @@ def route_config_proto(
         case "is_modifiable":
             logger.info("IS_MODIFIABLE")
             res = proto_base.ConfigResponse(session_id=config.session_id)
+            telemetry.report_config_get(config.operation.is_modifiable.keys)
             for key in config.operation.is_modifiable.keys:
                 pair = res.pairs.add()
                 pair.key = key
@@ -525,17 +531,8 @@ def set_snowflake_parameters(
                 snowpark_session.sql(
                     f"ALTER SESSION SET TIMEZONE = '{value}'"
                 ).collect()
-                set_jvm_timezone(value)
-                if hasattr(time, "tzset"):
-                    os.environ["TZ"] = value
-                    time.tzset()
             else:
                 snowpark_session.sql("ALTER SESSION UNSET TIMEZONE").collect()
-                reset_jvm_timezone_to_system_default()
-                if hasattr(time, "tzset") and False:
-                    if "TZ" in os.environ:
-                        del os.environ["TZ"]
-                    time.tzset()
         case "spark.sql.globalTempDatabase":
             if not value:
                 value = global_config.default_static_global_config.get(key)
@@ -588,3 +585,11 @@ def auto_uppercase_non_column_identifiers() -> bool:
     return session_config[
         "snowpark.connect.sql.identifiers.auto-uppercase"
     ].lower() in ("all", "all_except_columns")
+
+
+def parse_imports(session: snowpark.Session, imports: str | None) -> None:
+    if not imports:
+        return
+
+    for udf_import in imports.strip("[] ").split(","):
+        session.add_import(udf_import)

snowflake/snowpark_connect/error/error_utils.py

@@ -28,7 +28,9 @@ from pyspark.errors.exceptions.base import (
     PySparkException,
     PythonException,
     SparkRuntimeException,
+    UnsupportedOperationException,
 )
+from pyspark.errors.exceptions.connect import SparkConnectGrpcException
 from snowflake.core.exceptions import NotFoundError
 
 from snowflake.connector.errors import ProgrammingError
@@ -49,7 +51,9 @@ SPARK_PYTHON_TO_JAVA_EXCEPTION = {
     ArrayIndexOutOfBoundsException: "java.lang.ArrayIndexOutOfBoundsException",
     NumberFormatException: "java.lang.NumberFormatException",
     SparkRuntimeException: "org.apache.spark.SparkRuntimeException",
+    SparkConnectGrpcException: "pyspark.errors.exceptions.connect.SparkConnectGrpcException",
     PythonException: "org.apache.spark.api.python.PythonException",
+    UnsupportedOperationException: "java.lang.UnsupportedOperationException",
 }
 
 WINDOW_FUNCTION_ANALYSIS_EXCEPTION_SQL_ERROR_CODE = {1005, 2303}
@@ -68,6 +72,13 @@ init_multi_args_exception_pattern = (
 terminate_multi_args_exception_pattern = (
     r"terminate\(\) missing \d+ required positional argument"
 )
+snowpark_connect_exception_pattern = re.compile(
+    r"\[snowpark-connect-exception(?::(\w+))?\]\s*(.+?)'\s*is not recognized"
+)
+invalid_bit_pattern = re.compile(
+    r"Invalid bit position: \d+ exceeds the bit (?:upper|lower) limit",
+    re.IGNORECASE,
+)
 
 
 def contains_udtf_select(sql_string):
@@ -100,6 +111,22 @@ def _get_converted_known_sql_or_custom_exception(
         return SparkRuntimeException(
             message="Unexpected value for start in function slice: SQL array indices start at 1."
         )
+    invalid_bit = invalid_bit_pattern.search(msg)
+    if invalid_bit:
+        return IllegalArgumentException(message=invalid_bit.group(0))
+    match = snowpark_connect_exception_pattern.search(
+        ex.message if hasattr(ex, "message") else str(ex)
+    )
+    if match:
+        class_name = match.group(1)
+        message = match.group(2)
+        exception_class = (
+            globals().get(class_name, SparkConnectGrpcException)
+            if class_name
+            else SparkConnectGrpcException
+        )
+        return exception_class(message=message)
+
     if "select with no columns" in msg and contains_udtf_select(query):
         # We try our best to detect if the SQL string contains a UDTF call and the output schema is empty.
         return PythonException(message=f"[UDTF_RETURN_SCHEMA_MISMATCH] {ex.message}")
@@ -131,6 +158,11 @@ def _get_converted_known_sql_or_custom_exception(
                 message=f"[UDTF_EXEC_ERROR] User defined table function encountered an error in the terminate method: {ex.message}"
             )
 
+        if "failed to split string, provided pattern:" in msg:
+            return IllegalArgumentException(
+                message=f"Failed to split string using provided pattern. {ex.message}"
+            )
+
         if "100357" in msg and "wrong tuple size for returned value" in msg:
             return PythonException(
                 message=f"[UDTF_RETURN_SCHEMA_MISMATCH] The number of columns in the result does not match the specified schema. {ex.message}"

snowflake/snowpark_connect/error/exceptions.py

@@ -22,3 +22,7 @@ class MissingSchema(SnowparkConnectException):
         super().__init__(
             "No default schema found in session",
         )
+
+
+class MaxRetryExceeded(SnowparkConnectException):
+    ...

snowflake/snowpark_connect/expression/hybrid_column_map.py

@@ -0,0 +1,192 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+"""
+Hybrid column mapping for HAVING clause resolution.
+
+This module provides a special column mapping that can resolve expressions
+in the context of both the input DataFrame (for base columns) and the
+aggregated DataFrame (for aggregate expressions and aliases).
+"""
+
+from typing import Dict, List
+
+import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
+
+from snowflake import snowpark
+from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
+from snowflake.snowpark_connect.expression.typer import ExpressionTyper
+from snowflake.snowpark_connect.typed_column import TypedColumn
+
+
+class HybridColumnMap:
+    """
+    A column mapping that can resolve expressions in both input and aggregated contexts.
+
+    This is specifically designed for HAVING clause resolution where expressions may reference:
+    1. Base columns from the input DataFrame (to build new aggregates)
+    2. Existing aggregate expressions and their aliases
+    3. Grouping columns
+    """
+
+    def __init__(
+        self,
+        input_column_map: ColumnNameMap,
+        input_typer: ExpressionTyper,
+        aggregated_column_map: ColumnNameMap,
+        aggregated_typer: ExpressionTyper,
+        aggregate_expressions: List[expressions_proto.Expression],
+        grouping_expressions: List[expressions_proto.Expression],
+        aggregate_aliases: Dict[str, expressions_proto.Expression],
+    ) -> None:
+        self.input_column_map = input_column_map
+        self.input_typer = input_typer
+        self.aggregated_column_map = aggregated_column_map
+        self.aggregated_typer = aggregated_typer
+        self.aggregate_expressions = aggregate_expressions
+        self.grouping_expressions = grouping_expressions
+        self.aggregate_aliases = aggregate_aliases
+
+    def is_aggregate_function(self, exp: expressions_proto.Expression) -> bool:
+        """Check if an expression is an aggregate function."""
+        if exp.WhichOneof("expr_type") == "unresolved_function":
+            func_name = exp.unresolved_function.function_name.lower()
+            # Common aggregate functions - expand this list as needed
+            aggregate_functions = {
+                "avg",
+                "average",
+                "sum",
+                "count",
+                "min",
+                "max",
+                "stddev",
+                "stddev_pop",
+                "stddev_samp",
+                "variance",
+                "var_pop",
+                "var_samp",
+                "collect_list",
+                "collect_set",
+                "first",
+                "last",
+                "any_value",
+                "bool_and",
+                "bool_or",
+                "corr",
+                "covar_pop",
+                "covar_samp",
+                "kurtosis",
+                "skewness",
+                "percentile_cont",
+                "percentile_disc",
+                "approx_count_distinct",
+            }
+            return func_name in aggregate_functions
+        return False
+
+    def is_grouping_column(self, column_name: str) -> bool:
+        """Check if a column name refers to a grouping column."""
+        for group_exp in self.grouping_expressions:
+            if (
+                group_exp.WhichOneof("expr_type") == "unresolved_attribute"
+                and group_exp.unresolved_attribute.unparsed_identifier == column_name
+            ):
+                return True
+        return False
+
+    def resolve_expression(
+        self, exp: expressions_proto.Expression
+    ) -> tuple[list[str], TypedColumn]:
+        """
+        Resolve an expression in the hybrid context.
+
+        Strategy:
+        1. If it's an aggregate function -> create new aggregate using input context
+        2. If it's an alias to existing aggregate -> use aggregated context
+        3. If it's a grouping column -> try aggregated context first, fall back to input context
+           (handles exclude_grouping_columns=True case)
+        4. Otherwise -> try input context first, then aggregated context
+        """
+        from snowflake.snowpark_connect.expression.map_expression import map_expression
+
+        expr_type = exp.WhichOneof("expr_type")
+
+        # Handle aggregate functions - need to evaluate against input DataFrame
+        if self.is_aggregate_function(exp):
+            return map_expression(exp, self.input_column_map, self.input_typer)
+
+        # Handle column references
+        if expr_type == "unresolved_attribute":
+            column_name = exp.unresolved_attribute.unparsed_identifier
+
+            # Check if it's an alias to an existing aggregate expression
+            if column_name in self.aggregate_aliases:
+                # Use the aggregated context to get the alias
+                return map_expression(
+                    exp, self.aggregated_column_map, self.aggregated_typer
+                )
+
+            # Check if it's a grouping column
+            if self.is_grouping_column(column_name):
+                # Try aggregated context first (for cases where grouping columns are included)
+                try:
+                    return map_expression(
+                        exp, self.aggregated_column_map, self.aggregated_typer
+                    )
+                except Exception:
+                    # Fall back to input context if grouping columns were excluded
+                    # This handles the exclude_grouping_columns=True case
+                    return map_expression(exp, self.input_column_map, self.input_typer)
+
+            # Try input context first (for base columns used in new aggregates)
+            try:
+                return map_expression(exp, self.input_column_map, self.input_typer)
+            except Exception:
+                # Fall back to aggregated context
+                return map_expression(
+                    exp, self.aggregated_column_map, self.aggregated_typer
+                )
+
+        # For other expression types, try aggregated context first (likely references to computed values)
+        try:
+            return map_expression(
+                exp, self.aggregated_column_map, self.aggregated_typer
+            )
+        except Exception:
+            # Fall back to input context
+            return map_expression(exp, self.input_column_map, self.input_typer)
+
+
+def create_hybrid_column_map_for_having(
+    input_df: snowpark.DataFrame,
+    input_column_map: ColumnNameMap,
+    aggregated_df: snowpark.DataFrame,
+    aggregated_column_map: ColumnNameMap,
+    aggregate_expressions: List[expressions_proto.Expression],
+    grouping_expressions: List[expressions_proto.Expression],
+    spark_columns: List[str],
+    raw_aggregations: List[tuple[str, TypedColumn]],
+) -> HybridColumnMap:
+    """
+    Create a HybridColumnMap instance for HAVING clause resolution.
+    """
+    # Create typers for both contexts
+    input_typer = ExpressionTyper(input_df)
+    aggregated_typer = ExpressionTyper(aggregated_df)
+
+    # Build alias mapping from spark column names to aggregate expressions
+    aggregate_aliases = {}
+    for i, (spark_name, _) in enumerate(raw_aggregations):
+        if i < len(aggregate_expressions):
+            aggregate_aliases[spark_name] = aggregate_expressions[i]
+
+    return HybridColumnMap(
+        input_column_map=input_column_map,
+        input_typer=input_typer,
+        aggregated_column_map=aggregated_column_map,
+        aggregated_typer=aggregated_typer,
+        aggregate_expressions=aggregate_expressions,
+        grouping_expressions=grouping_expressions,
+        aggregate_aliases=aggregate_aliases,
+    )

snowflake/snowpark_connect/expression/literal.py

@@ -10,6 +10,7 @@ import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 from tzlocal import get_localzone
 
 from snowflake.snowpark_connect.config import global_config
+from snowflake.snowpark_connect.utils.context import get_is_evaluating_sql
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
@@ -47,17 +48,7 @@ def get_literal_field_and_name(literal: expressions_proto.Expression.Literal):
             ).date()
             return date, f"DATE '{date}'"
         case "timestamp" | "timestamp_ntz" as t:
-            # Note - Clients need to ensure local_timezone is the same as spark_sql_session_timeZone config.
-            # No need to apply timezone for lit datetime, because we set the TIMEZONE parameter in snowpark session,
-            # the snowflake backend would convert the lit datetime correctly. However, for returned column name, the
-            # timezone needs to be added. Pyspark has a weird behavior that datetime.datetime always gets converted
-            # to local timezone before printing according to spark_sql_session_timeZone setting. Haven't found
-            # official doc about it, but this behavior is based on my testings.
-            tz = (
-                ZoneInfo(global_config.spark_sql_session_timeZone)
-                if hasattr(global_config, "spark_sql_session_timeZone")
-                else get_localzone()
-            )
+            local_tz = get_localzone()
             if t == "timestamp":
                 microseconds = literal.timestamp
             else:
@@ -66,11 +57,17 @@ def get_literal_field_and_name(literal: expressions_proto.Expression.Literal):
                 microseconds // 1_000_000
             ) + datetime.timedelta(microseconds=microseconds % 1_000_000)
             tz_dt = datetime.datetime.fromtimestamp(
-                microseconds // 1_000_000, tz=tz
+                microseconds // 1_000_000, tz=local_tz
             ) + datetime.timedelta(microseconds=microseconds % 1_000_000)
             if t == "timestamp_ntz":
                 lit_dt = lit_dt.astimezone(datetime.timezone.utc)
                 tz_dt = tz_dt.astimezone(datetime.timezone.utc)
+            elif not get_is_evaluating_sql():
+                config_tz = global_config.spark_sql_session_timeZone
+                config_tz = ZoneInfo(config_tz) if config_tz else local_tz
+                tz_dt = tz_dt.astimezone(config_tz)
+                lit_dt = lit_dt.astimezone(local_tz)
+
             return lit_dt, f"{t.upper()} '{tz_dt.strftime('%Y-%m-%d %H:%M:%S')}'"
         case "day_time_interval":
             # TODO(SNOW-1920942): Snowflake SQL is missing an "interval" type.

snowflake/snowpark_connect/expression/map_cast.py

@@ -127,10 +127,11 @@ def map_cast(
         from_type = StringType()
     if isinstance(to_type, StringType):
         to_type = StringType()
+
+    # todo - verify if that's correct SNOW-2248680
     if isinstance(from_type, TimestampType):
         from_type = TimestampType()
-    if isinstance(to_type, TimestampType):
-        to_type = TimestampType()
+
     match (from_type, to_type):
         case (_, _) if (from_type == to_type):
             result_exp = col
@@ -185,6 +186,17 @@ def map_cast(
         case (DateType(), TimestampType()):
             result_exp = snowpark_fn.to_timestamp(col)
             result_exp = result_exp.cast(TimestampType(TimestampTimeZone.NTZ))
+        case (TimestampType() as f, TimestampType() as t) if f.tzinfo == t.tzinfo:
+            result_exp = col
+        case (
+            TimestampType(),
+            TimestampType() as t,
+        ) if t.tzinfo == TimestampTimeZone.NTZ:
+            zone = global_config.spark_sql_session_timeZone
+            result_exp = snowpark_fn.convert_timezone(snowpark_fn.lit(zone), col).cast(
+                TimestampType(TimestampTimeZone.NTZ)
+            )
+        # todo: verify if more support for LTZ and TZ is needed - SNOW-2248680
         case (TimestampType(), TimestampType()):
             result_exp = col
         case (_, TimestampType()) if isinstance(from_type, _NumericType):
@@ -259,8 +271,12 @@ def map_cast(
         case (_, _) if isinstance(from_type, _FractionalType) and isinstance(
             to_type, _IntegralType
         ):
-            result_exp = snowpark_fn.when(col < 0, snowpark_fn.ceil(col)).otherwise(
-                snowpark_fn.floor(col)
+            result_exp = (
+                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))
             )
             result_exp = result_exp.cast(to_type)
         case (StringType(), _) if (isinstance(to_type, _IntegralType)):

snowflake/snowpark_connect/expression/map_expression.py

@@ -11,6 +11,7 @@ import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
 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.expression import (
     map_extension,
@@ -190,9 +191,15 @@ def map_expression(
                 return [lit_name], TypedColumn(
                     snowpark_fn.lit(lit_value, return_type), lambda: [return_type]
                 )
+            result_exp = snowpark_fn.lit(lit_value)
+
+            if lit_type_str == "timestamp_ntz" and isinstance(
+                lit_value, datetime.datetime
+            ):
+                result_exp = result_exp.cast(TimestampType(TimestampTimeZone.NTZ))
 
             return [lit_name], TypedColumn(
-                snowpark_fn.lit(lit_value), lambda: [map_simple_types(lit_type_str)]
+                result_exp, lambda: [map_simple_types(lit_type_str)]
             )
         case "sort_order":
             child_name, child_column = map_single_column_expression(

snowflake/snowpark_connect/expression/map_udf.py

@@ -13,10 +13,7 @@ from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.type_mapping import proto_to_snowpark_type
 from snowflake.snowpark_connect.typed_column import TypedColumn
-from snowflake.snowpark_connect.utils.session import (
-    get_or_create_snowpark_session,
-    get_python_udxf_import_files,
-)
+from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
 from snowflake.snowpark_connect.utils.udf_helper import (
     SnowparkUDF,
     gen_input_types,
@@ -28,6 +25,9 @@ from snowflake.snowpark_connect.utils.udf_helper import (
 from snowflake.snowpark_connect.utils.udf_utils import (
     ProcessCommonInlineUserDefinedFunction,
 )
+from snowflake.snowpark_connect.utils.udxf_import_utils import (
+    get_python_udxf_import_files,
+)
 
 
 def process_udf_return_type(

snowflake/snowpark_connect/expression/map_unresolved_extract_value.py

@@ -5,6 +5,7 @@
 import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 
 import snowflake.snowpark.functions as snowpark_fn
+from snowflake.snowpark._internal.analyzer.expression import Literal
 from snowflake.snowpark.types import ArrayType, MapType, StructType, _IntegralType
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
 from snowflake.snowpark_connect.config import global_config
@@ -57,7 +58,8 @@ def map_unresolved_extract_value(
         extract_fn = snowpark_fn.get_ignore_case
     # Set index to a dummy value before we use it later in the ansi mode check.
     index = snowpark_fn.lit(1)
-    if _check_if_array_type(extract_typed_column, child_typed_column):
+    is_array = _check_if_array_type(extract_typed_column, child_typed_column)
+    if is_array:
         # Set all non-valid array indices to NULL.
         # This is done because both conditions of a CASE WHEN statement are executed regardless of if the condition is true or not.
         # Getting a negative index in Snowflake throws an error; thus, we convert all non-valid array indices to NULL before getting the index.
@@ -74,12 +76,37 @@ def map_unresolved_extract_value(
 
     spark_sql_ansi_enabled = global_config.spark_sql_ansi_enabled
 
-    if spark_sql_ansi_enabled and _check_if_array_type(
-        extract_typed_column, child_typed_column
-    ):
+    if spark_sql_ansi_enabled and is_array:
         result_exp = snowpark_fn.when(
             index.isNull(),
             child_typed_column.col.getItem("[snowpark_connect::INVALID_ARRAY_INDEX]"),
         ).otherwise(result_exp)
 
-    return spark_function_name, TypedColumn(result_exp, lambda: typer.type(result_exp))
+    def _get_extracted_value_type():
+        if is_array:
+            return [child_typed_column.typ.element_type]
+        elif isinstance(child_typed_column.typ, MapType):
+            return [child_typed_column.typ.value_type]
+        elif (
+            isinstance(child_typed_column.typ, StructType)
+            and isinstance(extract_typed_column.col._expr1, Literal)
+            and isinstance(extract_typed_column.col._expr1.value, str)
+        ):
+            struct = dict(
+                {
+                    (
+                        f.name
+                        if global_config.spark_sql_caseSensitive
+                        else f.name.lower(),
+                        f.datatype,
+                    )
+                    for f in child_typed_column.typ.fields
+                }
+            )
+            key = extract_typed_column.col._expr1.value
+            key = key if global_config.spark_sql_caseSensitive else key.lower()
+
+            return [struct[key]] if key in struct else typer.type(result_exp)
+        return typer.type(result_exp)
+
+    return spark_function_name, TypedColumn(result_exp, _get_extracted_value_type)