snowpark-connect 0.22.1__tar.gz → 0.24.0__tar.gz

{snowpark_connect-0.22.1/src/snowpark_connect.egg-info → snowpark_connect-0.24.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowpark-connect
-Version: 0.22.1
+Version: 0.24.0
 Summary: Snowpark Connect for Spark
 Author: Snowflake, Inc
 License: Apache License, Version 2.0
@@ -16,7 +16,7 @@ Requires-Dist: jpype1
 Requires-Dist: protobuf<5.0,>=4.25.3
 Requires-Dist: s3fs>=2025.3.0
 Requires-Dist: snowflake.core<2,>=1.0.5
-Requires-Dist: snowflake-snowpark-python[pandas]!=1.35.0,!=1.36.0,>=1.34.0
+Requires-Dist: snowflake-snowpark-python[pandas]<1.38.0,>=1.37.0
 Requires-Dist: sqlglot>=26.3.8
 Requires-Dist: jaydebeapi
 Requires-Dist: aiobotocore~=2.23.0

{snowpark_connect-0.22.1 → snowpark_connect-0.24.0}/setup.py

@@ -46,7 +46,7 @@ setup(
         "protobuf>=4.25.3,<5.0",
         "s3fs>=2025.3.0",  # prod-297255-inc0132291
         "snowflake.core>=1.0.5,<2",
-        "snowflake-snowpark-python[pandas]>=1.34.0,!=1.35.0,!=1.36.0",
+        "snowflake-snowpark-python[pandas]>=1.37.0,<1.38.0",
         "sqlglot>=26.3.8",
         "jaydebeapi",
         "aiobotocore~=2.23.0",

{snowpark_connect-0.22.1 → snowpark_connect-0.24.0}/src/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
@@ -533,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)

{snowpark_connect-0.22.1 → snowpark_connect-0.24.0}/src/snowflake/snowpark_connect/error/error_utils.py

@@ -75,6 +75,10 @@ terminate_multi_args_exception_pattern = (
 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):
@@ -107,6 +111,9 @@ 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)
     )

{snowpark_connect-0.22.1 → snowpark_connect-0.24.0}/src/snowflake/snowpark_connect/error/exceptions.py

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

snowpark_connect-0.24.0/src/snowflake/snowpark_connect/expression/function_defaults.py

@@ -0,0 +1,207 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+from dataclasses import dataclass
+from typing import Any
+
+import pyspark.sql.connect.proto.expressions_pb2 as expressions_pb2
+import pyspark.sql.connect.proto.types_pb2 as types_pb2
+
+
+@dataclass(frozen=True)
+class DefaultParameter:
+    """Represents a single default parameter for a function."""
+
+    name: str
+    value: Any
+
+
+@dataclass(frozen=True)
+class FunctionDefaults:
+    """Represents default parameter configuration for a function."""
+
+    total_args: int
+    defaults: list[DefaultParameter]
+
+
+# FUNCTION_DEFAULTS dictionary to hold operation name with default values.
+# This is required as non pyspark clients such as scala or sql won't send all the parameters.
+# We use this dict to inject the missing parameters before processing the unresolved function.
+FUNCTION_DEFAULTS: dict[str, FunctionDefaults] = {
+    "aes_decrypt": FunctionDefaults(
+        total_args=5,
+        defaults=[
+            DefaultParameter("mode", "GCM"),  # Spark SQL default: GCM
+            DefaultParameter("padding", "NONE"),  # Spark SQL default: NONE for GCM mode
+            DefaultParameter("aad", ""),  # Spark SQL default: empty string
+        ],
+    ),
+    "aes_encrypt": FunctionDefaults(
+        total_args=6,
+        defaults=[
+            DefaultParameter("mode", "GCM"),  # Spark SQL default: GCM
+            DefaultParameter("padding", "NONE"),  # Spark SQL default: NONE for GCM mode
+            DefaultParameter(
+                "iv", ""
+            ),  # Spark SQL default: empty string (random generated if not provided)
+            DefaultParameter("aad", ""),  # Spark SQL default: empty string
+        ],
+    ),
+    "approx_percentile": FunctionDefaults(
+        total_args=3,
+        defaults=[DefaultParameter("accuracy", 10000)],
+    ),
+    "bround": FunctionDefaults(
+        total_args=2,
+        defaults=[DefaultParameter("scale", 0)],
+    ),
+    "first": FunctionDefaults(
+        total_args=2,
+        defaults=[DefaultParameter("ignorenulls", False)],
+    ),
+    "lag": FunctionDefaults(
+        total_args=2,
+        defaults=[
+            DefaultParameter("offset", 1),
+        ],
+    ),
+    "last": FunctionDefaults(
+        total_args=2,
+        defaults=[DefaultParameter("ignorenulls", False)],
+    ),
+    "lead": FunctionDefaults(
+        total_args=3,
+        defaults=[DefaultParameter("offset", 1), DefaultParameter("default", None)],
+    ),
+    "locate": FunctionDefaults(
+        total_args=3,
+        defaults=[DefaultParameter("pos", 1)],
+    ),
+    "months_between": FunctionDefaults(
+        total_args=3,
+        defaults=[DefaultParameter("roundOff", True)],
+    ),
+    "nth_value": FunctionDefaults(
+        total_args=3,
+        defaults=[DefaultParameter("ignoreNulls", False)],
+    ),
+    "overlay": FunctionDefaults(
+        total_args=4,
+        defaults=[DefaultParameter("len", -1)],
+    ),
+    "percentile": FunctionDefaults(
+        total_args=3,
+        defaults=[DefaultParameter("frequency", 1)],
+    ),
+    "percentile_approx": FunctionDefaults(
+        total_args=3,
+        defaults=[DefaultParameter("accuracy", 10000)],
+    ),
+    "round": FunctionDefaults(
+        total_args=2,
+        defaults=[DefaultParameter("scale", 0)],
+    ),
+    "sentences": FunctionDefaults(
+        total_args=3,
+        defaults=[
+            DefaultParameter("language", ""),
+            DefaultParameter("country", ""),
+        ],
+    ),
+    "sort_array": FunctionDefaults(
+        total_args=2,
+        defaults=[DefaultParameter("asc", True)],
+    ),
+    "split": FunctionDefaults(
+        total_args=3,
+        defaults=[DefaultParameter("limit", -1)],
+    ),
+    "str_to_map": FunctionDefaults(
+        total_args=3,
+        defaults=[
+            DefaultParameter(
+                "pairDelim", ","
+            ),  # Spark SQL default: comma for splitting pairs
+            DefaultParameter(
+                "keyValueDelim", ":"
+            ),  # Spark SQL default: colon for splitting key/value
+        ],
+    ),
+    "try_aes_decrypt": FunctionDefaults(
+        total_args=5,
+        defaults=[
+            DefaultParameter("mode", "GCM"),  # Spark SQL default: GCM
+            DefaultParameter("padding", "NONE"),  # Spark SQL default: NONE for GCM mode
+            DefaultParameter("aad", ""),  # Spark SQL default: empty string
+        ],
+    ),
+}
+
+
+def _create_literal_expression(value: Any) -> expressions_pb2.Expression:
+    """Create a literal expression for the given value."""
+    expr = expressions_pb2.Expression()
+    if isinstance(value, bool):
+        expr.literal.boolean = value
+    elif isinstance(value, int):
+        expr.literal.integer = value
+    elif isinstance(value, str):
+        expr.literal.string = value
+    elif isinstance(value, float):
+        expr.literal.double = value
+    elif value is None:
+        null_type = types_pb2.DataType()
+        null_type.null.SetInParent()
+        expr.literal.null.CopyFrom(null_type)
+    else:
+        raise ValueError(f"Unsupported literal type: {value}")
+
+    return expr
+
+
+def inject_function_defaults(
+    unresolved_function: expressions_pb2.Expression.UnresolvedFunction,
+) -> bool:
+    """
+    Inject missing default parameters into an UnresolvedFunction protobuf.
+
+    Args:
+        unresolved_function: The protobuf UnresolvedFunction to modify
+
+    Returns:
+        bool: True if any defaults were injected, False otherwise
+    """
+    function_name = unresolved_function.function_name.lower()
+
+    if function_name not in FUNCTION_DEFAULTS:
+        return False
+
+    func_config = FUNCTION_DEFAULTS[function_name]
+    current_arg_count = len(unresolved_function.arguments)
+    total_args = func_config.total_args
+    defaults = func_config.defaults
+
+    if not defaults or current_arg_count >= total_args:
+        return False
+
+    # Calculate how many defaults to append
+    missing_arg_count = total_args - current_arg_count
+
+    # Check if any required params are missing.
+    if missing_arg_count > len(defaults):
+        raise ValueError(
+            f"Function '{function_name}' is missing required arguments. "
+            f"Expected {total_args} args, got {current_arg_count}, "
+            f"but only {len(defaults)} defaults are defined."
+        )
+
+    defaults_to_append = defaults[-missing_arg_count:]
+    injected = False
+
+    # Simply append the needed default values
+    for default_param in defaults_to_append:
+        default_expr = _create_literal_expression(default_param.value)
+        unresolved_function.arguments.append(default_expr)
+        injected = True
+
+    return injected

snowpark_connect-0.24.0/src/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,
+    )

{snowpark_connect-0.22.1 → snowpark_connect-0.24.0}/src/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.
@@ -84,6 +81,11 @@ def get_literal_field_and_name(literal: expressions_proto.Expression.Literal):
         case "decimal":
             # literal.decimal.precision & scale are ignored, as decimal.Decimal doesn't accept them
             return decimal.Decimal(literal.decimal.value), literal.decimal.value
+        case "array":
+            array_values, element_names = zip(
+                *(get_literal_field_and_name(e) for e in literal.array.elements)
+            )
+            return array_values, f"ARRAY({', '.join(element_names)})"
         case "null" | None:
             return None, "NULL"
         case other:

{snowpark_connect-0.22.1 → snowpark_connect-0.24.0}/src/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)):

{snowpark_connect-0.22.1 → snowpark_connect-0.24.0}/src/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,
@@ -26,7 +27,10 @@ from snowflake.snowpark_connect.expression.literal import get_literal_field_and_
 from snowflake.snowpark_connect.expression.map_cast import map_cast
 from snowflake.snowpark_connect.expression.map_sql_expression import map_sql_expr
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
-from snowflake.snowpark_connect.type_mapping import map_simple_types
+from snowflake.snowpark_connect.type_mapping import (
+    map_simple_types,
+    proto_to_snowpark_type,
+)
 from snowflake.snowpark_connect.typed_column import TypedColumn
 from snowflake.snowpark_connect.utils.context import (
     gen_sql_plan_id,
@@ -165,6 +169,12 @@ def map_expression(
                     lambda: [map_simple_types(lit_type_str)],
                 )
 
+            if lit_type_str == "array":
+                result_exp = snowpark_fn.lit(lit_value)
+                element_types = proto_to_snowpark_type(exp.literal.array.element_type)
+                array_type = snowpark.types.ArrayType(element_types)
+                return [lit_name], TypedColumn(result_exp, lambda: [array_type])
+
             # Decimal needs further processing to get the precision and scale properly.
             if lit_type_str == "decimal":
                 # Precision and scale are optional in the proto.
@@ -190,9 +200,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(

{snowpark_connect-0.22.1 → snowpark_connect-0.24.0}/src/snowflake/snowpark_connect/expression/map_extension.py

@@ -10,7 +10,10 @@ from snowflake.snowpark.types import BooleanType
 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
-from snowflake.snowpark_connect.utils.context import push_evaluating_sql_scope
+from snowflake.snowpark_connect.utils.context import (
+    push_evaluating_sql_scope,
+    push_outer_dataframe,
+)
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
@@ -52,12 +55,19 @@ def map_extension(
             return [name], typed_col
 
         case "subquery_expression":
+            from snowflake.snowpark_connect.dataframe_container import (
+                DataFrameContainer,
+            )
             from snowflake.snowpark_connect.expression.map_expression import (
                 map_expression,
             )
             from snowflake.snowpark_connect.relation.map_relation import map_relation
 
-            with push_evaluating_sql_scope():
+            current_outer_df = DataFrameContainer(
+                dataframe=typer.df, column_map=column_mapping
+            )
+
+            with push_evaluating_sql_scope(), push_outer_dataframe(current_outer_df):
                 df_container = map_relation(extension.subquery_expression.input)
                 df = df_container.dataframe