snowpark-connect 0.26.0__tar.gz → 0.27.0__tar.gz

{snowpark_connect-0.26.0/src/snowpark_connect.egg-info → snowpark_connect-0.27.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowpark-connect
-Version: 0.26.0
+Version: 0.27.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.38.0,>=1.37.0
+Requires-Dist: snowflake-snowpark-python[pandas]<1.39.0,==1.38.0
 Requires-Dist: sqlglot>=26.3.8
 Requires-Dist: jaydebeapi
 Requires-Dist: aiobotocore~=2.23.0

{snowpark_connect-0.26.0 → snowpark_connect-0.27.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.37.0,<1.38.0",
+        "snowflake-snowpark-python[pandas]==1.38.0,<1.39.0",
         "sqlglot>=26.3.8",
         "jaydebeapi",
         "aiobotocore~=2.23.0",

{snowpark_connect-0.26.0 → snowpark_connect-0.27.0}/src/snowflake/snowpark_connect/config.py

@@ -171,9 +171,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
         ),

{snowpark_connect-0.26.0 → snowpark_connect-0.27.0}/src/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,6 +35,7 @@ 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,
 )
@@ -271,6 +273,19 @@ 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:
+                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."
+                )
+
             col = snowpark_fn.Column(
                 UnresolvedAttribute(exp.unresolved_named_lambda_variable.name_parts[0])
             )

{snowpark_connect-0.26.0 → snowpark_connect-0.27.0}/src/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 (
@@ -356,6 +357,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.'

{snowpark_connect-0.26.0 → snowpark_connect-0.27.0}/src/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],
             )
@@ -8725,7 +8731,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 +9917,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 +9928,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:

{snowpark_connect-0.26.0 → snowpark_connect-0.27.0}/src/snowflake/snowpark_connect/relation/map_extension.py

@@ -23,6 +23,7 @@ from snowflake.snowpark_connect.relation.map_relation import map_relation
 from snowflake.snowpark_connect.typed_column import TypedColumn
 from snowflake.snowpark_connect.utils.context import (
     get_sql_aggregate_function_count,
+    not_resolving_fun_args,
     push_outer_dataframe,
     set_current_grouping_columns,
 )
@@ -335,14 +336,15 @@ def map_aggregate(
     typer = ExpressionTyper(input_df)
 
     def _map_column(exp: expression_proto.Expression) -> tuple[str, TypedColumn]:
-        new_names, snowpark_column = map_expression(
-            exp, input_container.column_map, typer
-        )
-        if len(new_names) != 1:
-            raise SnowparkConnectNotImplementedError(
-                "Multi-column aggregate expressions are not supported"
+        with not_resolving_fun_args():
+            new_names, snowpark_column = map_expression(
+                exp, input_container.column_map, typer
             )
-        return new_names[0], snowpark_column
+            if len(new_names) != 1:
+                raise SnowparkConnectNotImplementedError(
+                    "Multi-column aggregate expressions are not supported"
+                )
+            return new_names[0], snowpark_column
 
     raw_groupings: list[tuple[str, TypedColumn]] = []
     raw_aggregations: list[tuple[str, TypedColumn]] = []

snowpark_connect-0.27.0/src/snowflake/snowpark_connect/relation/map_map_partitions.py

@@ -0,0 +1,117 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import pyspark.sql.connect.proto.relations_pb2 as relation_proto
+from pyspark.sql.connect.proto.expressions_pb2 import CommonInlineUserDefinedFunction
+
+import snowflake.snowpark.functions as snowpark_fn
+from snowflake import snowpark
+from snowflake.snowpark.types import StructType
+from snowflake.snowpark_connect.constants import MAP_IN_ARROW_EVAL_TYPE
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
+from snowflake.snowpark_connect.relation.map_relation import map_relation
+from snowflake.snowpark_connect.type_mapping import proto_to_snowpark_type
+from snowflake.snowpark_connect.utils.context import map_partitions_depth
+from snowflake.snowpark_connect.utils.pandas_udtf_utils import (
+    create_pandas_udtf,
+    create_pandas_udtf_with_arrow,
+)
+from snowflake.snowpark_connect.utils.udf_helper import udf_check
+from snowflake.snowpark_connect.utils.udtf_helper import (
+    create_pandas_udtf_in_sproc,
+    require_creating_udtf_in_sproc,
+)
+
+
+def map_map_partitions(
+    rel: relation_proto.Relation,
+) -> DataFrameContainer:
+    """
+    Map a function over the partitions of the input DataFrame.
+
+    This is a simple wrapper around the `mapInPandas` method in Snowpark.
+    """
+    input_container = map_relation(rel.map_partitions.input)
+    udf_proto = rel.map_partitions.func
+    udf_check(udf_proto)
+
+    return _map_with_pandas_udtf(input_container, udf_proto)
+
+
+def _call_udtf(
+    udtf_name: str, input_df: snowpark.DataFrame, return_type: StructType | None = None
+) -> snowpark.DataFrame:
+    # Add a dummy column with random 1-10 values for partitioning
+    input_df_with_dummy = input_df.withColumn(
+        "_DUMMY_PARTITION_KEY",
+        (
+            snowpark_fn.uniform(
+                snowpark_fn.lit(1), snowpark_fn.lit(10), snowpark_fn.random()
+            )
+            * 10
+        ).cast("int"),
+    )
+
+    udtf_columns = input_df.columns + [snowpark_fn.col("_DUMMY_PARTITION_KEY")]
+
+    tfc = snowpark_fn.call_table_function(udtf_name, *udtf_columns).over(
+        partition_by=[snowpark_fn.col("_DUMMY_PARTITION_KEY")]
+    )
+
+    # Use map_partitions_depth only when mapping non nested map_partitions
+    # When mapping chained functions additional column casting is necessary
+    if map_partitions_depth() == 1:
+        result_df_with_dummy = input_df_with_dummy.join_table_function(tfc)
+    else:
+        result_df_with_dummy = input_df_with_dummy.select(tfc)
+
+    output_cols = [field.name for field in return_type.fields]
+
+    # Only return the output columns.
+    result_df = result_df_with_dummy.select(*output_cols)
+
+    return DataFrameContainer.create_with_column_mapping(
+        dataframe=result_df,
+        spark_column_names=output_cols,
+        snowpark_column_names=output_cols,
+        snowpark_column_types=[field.datatype for field in return_type.fields],
+    )
+
+
+def _map_with_pandas_udtf(
+    input_df_container: DataFrameContainer,
+    udf_proto: CommonInlineUserDefinedFunction,
+) -> snowpark.DataFrame:
+    """
+    Handle mapInArrow using pandas_udtf for partition-level Arrow processing.
+    """
+    input_df = input_df_container.dataframe
+    input_schema = input_df.schema
+    spark_column_names = input_df_container.column_map.get_spark_columns()
+    return_type = proto_to_snowpark_type(
+        udf_proto.python_udf.output_type
+        if udf_proto.WhichOneof("function") == "python_udf"
+        else udf_proto.scalar_scala_udf.outputType
+    )
+
+    # Check if this is mapInArrow (eval_type == 207)
+    map_in_arrow = (
+        udf_proto.WhichOneof("function") == "python_udf"
+        and udf_proto.python_udf.eval_type == MAP_IN_ARROW_EVAL_TYPE
+    )
+    if require_creating_udtf_in_sproc(udf_proto):
+        udtf_name = create_pandas_udtf_in_sproc(
+            udf_proto, spark_column_names, input_schema, return_type
+        )
+    else:
+        if map_in_arrow:
+            map_udtf = create_pandas_udtf_with_arrow(
+                udf_proto, spark_column_names, input_schema, return_type
+            )
+        else:
+            map_udtf = create_pandas_udtf(
+                udf_proto, spark_column_names, input_schema, return_type
+            )
+        udtf_name = map_udtf.name
+    return _call_udtf(udtf_name, input_df, return_type)

{snowpark_connect-0.26.0 → snowpark_connect-0.27.0}/src/snowflake/snowpark_connect/relation/map_relation.py

@@ -15,6 +15,7 @@ from snowflake.snowpark_connect.utils.cache import (
 from snowflake.snowpark_connect.utils.context import (
     get_plan_id_map,
     get_session_id,
+    push_map_partitions,
     push_operation_scope,
     set_is_aggregate_function,
     set_plan_id_map,
@@ -180,7 +181,8 @@ def map_relation(
                     )
                 return cached_df
             case "map_partitions":
-                result = map_map_partitions.map_map_partitions(rel)
+                with push_map_partitions():
+                    result = map_map_partitions.map_map_partitions(rel)
             case "offset":
                 result = map_row_ops.map_offset(rel)
             case "project":

{snowpark_connect-0.26.0 → snowpark_connect-0.27.0}/src/snowflake/snowpark_connect/relation/map_show_string.py

@@ -12,6 +12,7 @@ from snowflake.snowpark._internal.analyzer import analyzer_utils
 from snowflake.snowpark.functions import col
 from snowflake.snowpark.types import DateType, StringType, StructField, StructType
 from snowflake.snowpark_connect.column_name_handler import set_schema_getter
+from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
 from snowflake.snowpark_connect.relation.map_relation import map_relation
 
@@ -33,6 +34,7 @@ def map_show_string(rel: relation_proto.Relation) -> pandas.DataFrame:
         truncate=rel.show_string.truncate,
         vertical=rel.show_string.vertical,
         _spark_column_names=input_df_container.column_map.get_spark_columns(),
+        _spark_session_tz=global_config.spark_sql_session_timeZone,
     )
     return pandas.DataFrame({"show_string": [show_string]})
 

{snowpark_connect-0.26.0 → snowpark_connect-0.27.0}/src/snowflake/snowpark_connect/relation/map_sql.py

@@ -56,6 +56,7 @@ from snowflake.snowpark_connect.utils.context import (
     _accessing_temp_object,
     gen_sql_plan_id,
     get_session_id,
+    get_sql_plan,
     push_evaluating_sql_scope,
     push_sql_scope,
     set_sql_args,
@@ -542,6 +543,7 @@ def map_sql_to_pandas_df(
                 rows = session.sql(f"DESCRIBE TABLE {name}").collect()
             case "DescribeNamespace":
                 name = get_relation_identifier_name(logical_plan.namespace(), True)
+                name = change_default_to_public(name)
                 rows = session.sql(f"DESCRIBE SCHEMA {name}").collect()
                 if not rows:
                     rows = None
@@ -793,6 +795,7 @@ def map_sql_to_pandas_df(
             case "SetCatalogAndNamespace":
                 # TODO: add catalog setting here
                 name = get_relation_identifier_name(logical_plan.child(), True)
+                name = change_default_to_public(name)
                 session.sql(f"USE SCHEMA {name}").collect()
             case "SetCommand":
                 kv_result_tuple = logical_plan.kv().get()
@@ -801,6 +804,7 @@ def map_sql_to_pandas_df(
                 set_config_param(get_session_id(), key, val, session)
             case "SetNamespaceCommand":
                 name = _spark_to_snowflake(logical_plan.namespace())
+                name = change_default_to_public(name)
                 session.sql(f"USE SCHEMA {name}").collect()
             case "SetNamespaceLocation" | "SetNamespaceProperties":
                 raise SnowparkConnectNotImplementedError(
@@ -997,6 +1001,20 @@ def get_sql_passthrough() -> bool:
     return get_boolean_session_config_param("snowpark.connect.sql.passthrough")
 
 
+def change_default_to_public(name: str) -> str:
+    """
+    Change the namespace to PUBLIC when given name is DEFAULT
+    :param name: Given namespace
+    :return: if name is DEFAULT return PUBLIC otherwise name
+    """
+    if name.startswith('"'):
+        if name.upper() == '"DEFAULT"':
+            return name.replace("DEFAULT", "PUBLIC")
+    elif name.upper() == "DEFAULT":
+        return "PUBLIC"
+    return name
+
+
 def map_sql(
     rel: relation_proto.Relation,
 ) -> DataFrameContainer:
@@ -1008,7 +1026,6 @@ def map_sql(
     In passthough mode as True, SAS calls session.sql() and not calling Spark Parser.
     This is to mitigate any issue not covered by spark logical plan to protobuf conversion.
     """
-
     snowpark_connect_sql_passthrough = get_sql_passthrough()
 
     if not snowpark_connect_sql_passthrough:
@@ -1353,6 +1370,7 @@ def map_logical_plan_relation(
                     left_input=map_logical_plan_relation(children[0]),
                     right_input=map_logical_plan_relation(children[1]),
                     set_op_type=relation_proto.SetOperation.SET_OP_TYPE_UNION,
+                    is_all=True,
                     by_name=rel.byName(),
                     allow_missing_columns=rel.allowMissingCol(),
                 )
@@ -1701,7 +1719,50 @@ def map_logical_plan_relation(
                     _window_specs.get()[key] = window_spec
                 proto = map_logical_plan_relation(rel.child())
         case "Generate":
-            input_relation = map_logical_plan_relation(rel.child())
+            # Generate creates a nested Project relation (see lines 1785-1790) without
+            # setting its plan_id field. When this Project is later processed by map_project
+            # (map_column_ops.py), it uses rel.common.plan_id which defaults to 0 for unset
+            # protobuf fields. This means all columns from the Generate operation (both exploded
+            # columns and passthrough columns) will have plan_id=0 in their names.
+            #
+            # If Generate's child is a SubqueryAlias whose inner relation was processed
+            # with a non-zero plan_id, there will be a mismatch between:
+            # - The columns referenced in the Project (expecting plan_id from SubqueryAlias's child)
+            # - The actual column names created by Generate's Project (using plan_id=0)
+
+            # Therefore, when Generate has a SubqueryAlias child, we explicitly process the inner
+            # relation with plan_id=0 to match what Generate's Project will use. This only applies when
+            # the immediate child of Generate is a SubqueryAlias and preserves existing registrations (like CTEs),
+            # so it won't affect other patterns.
+
+            child_class = str(rel.child().getClass().getSimpleName())
+
+            if child_class == "SubqueryAlias":
+                alias = str(rel.child().alias())
+
+                # Check if this alias was already registered during initial SQL parsing
+                existing_plan_id = get_sql_plan(alias)
+
+                if existing_plan_id is not None:
+                    # Use the existing plan_id to maintain consistency with prior registration
+                    used_plan_id = existing_plan_id
+                else:
+                    # Use plan_id=0 to match what the nested Project will use (protobuf default)
+                    used_plan_id = 0
+                    set_sql_plan_name(alias, used_plan_id)
+
+                # Process the inner child with the determined plan_id
+                inner_child = map_logical_plan_relation(
+                    rel.child().child(), plan_id=used_plan_id
+                )
+                input_relation = relation_proto.Relation(
+                    subquery_alias=relation_proto.SubqueryAlias(
+                        input=inner_child,
+                        alias=alias,
+                    )
+                )
+            else:
+                input_relation = map_logical_plan_relation(rel.child())
             generator_output_list = as_java_list(rel.generatorOutput())
             generator_output_list_expressions = [
                 map_logical_plan_expression(e) for e in generator_output_list

{snowpark_connect-0.26.0 → snowpark_connect-0.27.0}/src/snowflake/snowpark_connect/relation/write/map_write.py

@@ -51,6 +51,7 @@ from snowflake.snowpark_connect.utils.identifiers import (
     split_fully_qualified_spark_name,
 )
 from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
+from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
     telemetry,
@@ -160,6 +161,29 @@ def map_write(request: proto_base.ExecutePlanRequest):
             compression_option = write_op.options.get("compression", "none")
 
             # Generate Spark-compatible filename or prefix
+            # we need a random prefix to support "append" mode
+            # otherwise copy into with overwrite=False will fail if the file already exists
+            overwrite = (
+                write_op.mode
+                == commands_proto.WriteOperation.SaveMode.SAVE_MODE_OVERWRITE
+            )
+
+            if overwrite:
+                try:
+                    path_after_stage = (
+                        write_path.split("/", 1)[1] if "/" in write_path else ""
+                    )
+                    if not path_after_stage or path_after_stage == "/":
+                        logger.warning(
+                            f"Skipping REMOVE for root path {write_path} - too broad scope"
+                        )
+                    else:
+                        remove_command = f"REMOVE {write_path}/"
+                        session.sql(remove_command).collect()
+                        logger.info(f"Successfully cleared directory: {write_path}")
+                except Exception as e:
+                    logger.warning(f"Could not clear directory {write_path}: {e}")
+
             if should_write_to_single_file:
                 # Single file: generate complete filename with extension
                 spark_filename = generate_spark_compatible_filename(
@@ -178,10 +202,6 @@ def map_write(request: proto_base.ExecutePlanRequest):
                     format_ext="",  # No extension for prefix
                 )
                 temp_file_prefix_on_stage = f"{write_path}/{spark_filename_prefix}"
-            overwrite = (
-                write_op.mode
-                == commands_proto.WriteOperation.SaveMode.SAVE_MODE_OVERWRITE
-            )
 
             default_compression = "NONE" if write_op.source != "parquet" else "snappy"
             compression = write_op.options.get(

{snowpark_connect-0.26.0 → snowpark_connect-0.27.0}/src/snowflake/snowpark_connect/server.py

@@ -531,7 +531,10 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
             if name.endswith(".class"):
                 # name is <dir>/<package>/<class_name>
                 # we don't need the dir name, but require the package, so only remove dir
-                class_files[name.split("/", 1)[-1]] = filepath
+                if os.name != "nt":
+                    class_files[name.split("/", 1)[-1]] = filepath
+                else:
+                    class_files[name.split("\\", 1)[-1]] = filepath
                 continue
             session.file.put(
                 filepath,

{snowpark_connect-0.26.0 → snowpark_connect-0.27.0}/src/snowflake/snowpark_connect/utils/artifacts.py

@@ -39,7 +39,7 @@ def write_temporary_artifact(
     if os.name != "nt":
         filepath = f"/tmp/sas-{session.session_id}/{name}"
     else:
-        filepath = f"{tempfile.gettempdir()}/sas-{session.session_id}/{name}"
+        filepath = f"{tempfile.gettempdir()}\\sas-{session.session_id}\\{name}"
     # The name comes to us as a path (e.g. cache/<name>), so we need to create
     # the parent directory if it doesn't exist to avoid errors during writing.
     pathlib.Path(filepath).parent.mkdir(parents=True, exist_ok=True)
@@ -55,11 +55,10 @@ def write_class_files_to_stage(
 ) -> None:
     if os.name != "nt":
         filepath = f"/tmp/sas-{session.session_id}"
+        jar_name = f'{filepath}/{hashlib.sha256(str(files).encode("utf-8")).hexdigest()[:10]}.jar'
     else:
-        filepath = f"{tempfile.gettempdir()}/sas-{session.session_id}"
-    jar_name = (
-        f'{filepath}/{hashlib.sha256(str(files).encode("utf-8")).hexdigest()[:10]}.jar'
-    )
+        filepath = f"{tempfile.gettempdir()}\\sas-{session.session_id}"
+        jar_name = f'{filepath}\\{hashlib.sha256(str(files).encode("utf-8")).hexdigest()[:10]}.jar'
     with zipfile.ZipFile(jar_name, "w", zipfile.ZIP_DEFLATED) as jar:
         for name, path in files.items():
             jar.write(path, name)

{snowpark_connect-0.26.0 → snowpark_connect-0.27.0}/src/snowflake/snowpark_connect/utils/context.py

@@ -30,6 +30,9 @@ _sql_aggregate_function_count = ContextVar[int](
     "_contains_aggregate_function", default=0
 )
 
+# Context for parsing map_partitions
+_map_partitions_stack = ContextVar[int]("_map_partitions_stack", default=0)
+
 # We have to generate our own plan IDs that are different from Spark's.
 # Spark plan IDs start at 0, so pick a "big enough" number to avoid overlaps.
 _STARTING_SQL_PLAN_ID = 0x80000000
@@ -49,6 +52,7 @@ _spark_client_type_regex = re.compile(r"spark/(?P<spark_version>\d+\.\d+\.\d+)")
 _current_operation = ContextVar[str]("_current_operation", default="default")
 _resolving_fun_args = ContextVar[bool]("_resolving_fun_args", default=False)
 _resolving_lambda_fun = ContextVar[bool]("_resolving_lambdas", default=False)
+_current_lambda_params = ContextVar[list[str]]("_current_lambda_params", default=[])
 
 _is_window_enabled = ContextVar[bool]("_is_window_enabled", default=False)
 _is_in_pivot = ContextVar[bool]("_is_in_pivot", default=False)
@@ -206,6 +210,16 @@ def push_evaluating_join_condition(join_type, left_keys, right_keys):
         _is_evaluating_join_condition.set(prev)
 
 
+@contextmanager
+def push_map_partitions():
+    _map_partitions_stack.set(_map_partitions_stack.get() + 1)
+    yield
+
+
+def map_partitions_depth() -> int:
+    return _map_partitions_stack.get()
+
+
 @contextmanager
 def push_sql_scope():
     """
@@ -238,16 +252,21 @@ def push_operation_scope(operation: str):
 
 
 @contextmanager
-def resolving_lambda_function():
+def resolving_lambda_function(param_names: list[str] = None):
     """
     Context manager that sets a flag indicating lambda function is being resolved.
+    Also tracks the lambda parameter names for validation.
     """
     prev = _resolving_lambda_fun.get()
+    prev_params = _current_lambda_params.get()
     try:
         _resolving_lambda_fun.set(True)
+        if param_names is not None:
+            _current_lambda_params.set(param_names)
         yield
     finally:
         _resolving_lambda_fun.set(prev)
+        _current_lambda_params.set(prev_params)
 
 
 def is_lambda_being_resolved() -> bool:
@@ -257,6 +276,13 @@ def is_lambda_being_resolved() -> bool:
     return _resolving_lambda_fun.get()
 
 
+def get_current_lambda_params() -> list[str]:
+    """
+    Returns the current lambda parameter names.
+    """
+    return _current_lambda_params.get()
+
+
 @contextmanager
 def resolving_fun_args():
     """
@@ -270,6 +296,19 @@ def resolving_fun_args():
         _resolving_fun_args.set(prev)
 
 
+@contextmanager
+def not_resolving_fun_args():
+    """
+    Context manager that sets a flag indicating function arguments are *not* being resolved.
+    """
+    prev = _resolving_fun_args.get()
+    try:
+        _resolving_fun_args.set(False)
+        yield
+    finally:
+        _resolving_fun_args.set(prev)
+
+
 def is_function_argument_being_resolved() -> bool:
     """
     Returns True if function arguments are being resolved.
@@ -350,6 +389,7 @@ def clear_context_data() -> None:
 
     _next_sql_plan_id.set(_STARTING_SQL_PLAN_ID)
     _sql_plan_name_map.set({})
+    _map_partitions_stack.set(0)
     _sql_aggregate_function_count.set(0)
     _sql_named_args.set({})
     _sql_pos_args.set({})