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

snowflake/snowpark_connect/config.py

@@ -22,6 +22,9 @@ from snowflake.snowpark.exceptions import SnowparkSQLException
 from snowflake.snowpark.types import TimestampTimeZone, TimestampType
 from snowflake.snowpark_connect.utils.concurrent import SynchronizedDict
 from snowflake.snowpark_connect.utils.context import get_session_id
+from snowflake.snowpark_connect.utils.external_udxf_cache import (
+    clear_external_udxf_cache,
+)
 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 (
@@ -136,6 +139,9 @@ class GlobalConfig:
         "spark.sql.parser.quotedRegexColumnNames": "false",
         # custom configs
         "snowpark.connect.version": ".".join(map(str, sas_version)),
+        # Control whether repartition(n) on a DataFrame forces splitting into n files during writes
+        # This matches spark behavior more closely, but introduces overhead.
+        "snowflake.repartition.for.writes": "false",
     }
 
     boolean_config_list = [
@@ -148,6 +154,7 @@ class GlobalConfig:
         "spark.sql.legacy.allowHashOnMapType",
         "spark.Catalog.databaseFilterInformationSchema",
         "spark.sql.parser.quotedRegexColumnNames",
+        "snowflake.repartition.for.writes",
     ]
 
     int_config_list = [
@@ -164,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
         ),
@@ -592,6 +596,9 @@ def parse_imports(session: snowpark.Session, imports: str | None) -> None:
     if not imports:
         return
 
+    # UDF needs to be recreated to include new imports
+    clear_external_udxf_cache(session)
+
     for udf_import in imports.strip("[] ").split(","):
         session.add_import(udf_import)
 

snowflake/snowpark_connect/dataframe_container.py

@@ -29,6 +29,7 @@ class DataFrameContainer:
         table_name: str | None = None,
         alias: str | None = None,
         cached_schema_getter: Callable[[], StructType] | None = None,
+        partition_hint: int | None = None,
     ) -> None:
         """
         Initialize a new DataFrameContainer.
@@ -39,11 +40,13 @@ class DataFrameContainer:
             table_name: Optional table name for the DataFrame
             alias: Optional alias for the DataFrame
             cached_schema_getter: Optional function to get cached schema
+            partition_hint: Optional partition count from repartition() operations
         """
         self._dataframe = dataframe
         self._column_map = self._create_default_column_map(column_map)
         self._table_name = table_name
         self._alias = alias
+        self._partition_hint = partition_hint
 
         if cached_schema_getter is not None:
             self._apply_cached_schema_getter(cached_schema_getter)
@@ -62,6 +65,7 @@ class DataFrameContainer:
         table_name: str | None = None,
         alias: str | None = None,
         cached_schema_getter: Callable[[], StructType] | None = None,
+        partition_hint: int | None = None,
     ) -> DataFrameContainer:
         """
         Create a new container with complete column mapping configuration.
@@ -78,6 +82,7 @@ class DataFrameContainer:
             table_name: Optional table name
             alias: Optional alias
             cached_schema_getter: Optional function to get cached schema
+            partition_hint: Optional partition count from repartition() operations
 
         Returns:
             A new DataFrameContainer instance
@@ -123,6 +128,7 @@ class DataFrameContainer:
             table_name=table_name,
             alias=alias,
             cached_schema_getter=final_schema_getter,
+            partition_hint=partition_hint,
         )
 
     @property
@@ -163,6 +169,16 @@ class DataFrameContainer:
         """Set the alias name."""
         self._alias = value
 
+    @property
+    def partition_hint(self) -> int | None:
+        """Get the partition hint count."""
+        return self._partition_hint
+
+    @partition_hint.setter
+    def partition_hint(self, value: int | None) -> None:
+        """Set the partition hint count."""
+        self._partition_hint = value
+
     def _create_default_column_map(
         self, column_map: ColumnNameMap | None
     ) -> ColumnNameMap:

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])
             )

snowflake/snowpark_connect/expression/map_udf.py

@@ -13,6 +13,10 @@ 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.external_udxf_cache import (
+    cache_external_udf,
+    get_external_udf_from_cache,
+)
 from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
 from snowflake.snowpark_connect.utils.udf_helper import (
     SnowparkUDF,
@@ -30,6 +34,39 @@ from snowflake.snowpark_connect.utils.udxf_import_utils import (
 )
 
 
+def cache_external_udf_wrapper(from_register_udf: bool):
+    def outer_wrapper(wrapper_func):
+        def wrapper(
+            udf_proto: expressions_proto.CommonInlineUserDefinedFunction,
+        ) -> SnowparkUDF | None:
+            udf_hash = hash(str(udf_proto))
+            cached_udf = get_external_udf_from_cache(udf_hash)
+
+            if cached_udf:
+                session = get_or_create_snowpark_session()
+                function_type = udf_proto.WhichOneof("function")
+                # TODO: Align this with SNOW-2316798 after merge
+                match function_type:
+                    case "scalar_scala_udf":
+                        session._udfs[cached_udf.name] = cached_udf
+                    case "python_udf" if from_register_udf:
+                        session._udfs[udf_proto.function_name.lower()] = cached_udf
+                    case "python_udf":
+                        pass
+                    case _:
+                        raise ValueError(f"Unsupported UDF type: {function_type}")
+
+                return cached_udf
+
+            snowpark_udf = wrapper_func(udf_proto)
+            cache_external_udf(udf_hash, snowpark_udf)
+            return snowpark_udf
+
+        return wrapper
+
+    return outer_wrapper
+
+
 def process_udf_return_type(
     return_type: types_proto.DataType,
 ) -> tuple[snowpark.types.DataType, snowpark.types.DataType]:
@@ -49,6 +86,7 @@ def process_udf_return_type(
     return original_snowpark_type, original_snowpark_type
 
 
+@cache_external_udf_wrapper(from_register_udf=True)
 def register_udf(
     udf_proto: expressions_proto.CommonInlineUserDefinedFunction,
 ) -> SnowparkUDF:
@@ -84,12 +122,10 @@ def register_udf(
             return_type=udf._return_type,
             original_return_type=original_return_type,
         )
-        # the create udf does register the udf but this seems to be for the client side check
-        # TODO: check if this is needed
+        session._udfs[udf_proto.function_name.lower()] = udf
+        # scala udfs can be also accessed using `udf.name`
         if udf_processor._function_type == "scalar_scala_udf":
             session._udfs[udf.name] = udf
-        else:
-            session._udfs[udf_proto.function_name.lower()] = udf
         return udf
 
 
@@ -114,29 +150,34 @@ def map_common_inline_user_defined_udf(
                 udf_proto.scalar_scala_udf.outputType
             )
 
-    session = get_or_create_snowpark_session()
-    kwargs = {
-        "common_inline_user_defined_function": udf_proto,
-        "input_types": input_types,
-        "called_from": "map_common_inline_user_defined_udf",
-        "return_type": processed_return_type,
-        "udf_packages": global_config.get("snowpark.connect.udf.packages", ""),
-        "udf_imports": get_python_udxf_import_files(session),
-        "original_return_type": original_return_type,
-    }
-    if require_creating_udf_in_sproc(udf_proto):
-        snowpark_udf = process_udf_in_sproc(**kwargs)
-    else:
-        udf_processor = ProcessCommonInlineUserDefinedFunction(**kwargs)
-        udf = udf_processor.create_udf()
-        snowpark_udf = SnowparkUDF(
-            name=udf.name,
-            input_types=udf._input_types,
-            return_type=udf._return_type,
-            original_return_type=original_return_type,
-        )
-        if udf_processor._function_type == "scalar_scala_udf":
-            session._udfs[udf.name] = snowpark_udf
+    @cache_external_udf_wrapper(from_register_udf=False)
+    def get_snowpark_udf(
+        udf_proto: expressions_proto.CommonInlineUserDefinedFunction,
+    ) -> SnowparkUDF:
+        session = get_or_create_snowpark_session()
+        kwargs = {
+            "common_inline_user_defined_function": udf_proto,
+            "input_types": input_types,
+            "called_from": "map_common_inline_user_defined_udf",
+            "return_type": processed_return_type,
+            "udf_packages": global_config.get("snowpark.connect.udf.packages", ""),
+            "udf_imports": get_python_udxf_import_files(session),
+            "original_return_type": original_return_type,
+        }
+        if require_creating_udf_in_sproc(udf_proto):
+            snowpark_udf = process_udf_in_sproc(**kwargs)
+        else:
+            udf_processor = ProcessCommonInlineUserDefinedFunction(**kwargs)
+            udf = udf_processor.create_udf()
+            snowpark_udf = SnowparkUDF(
+                name=udf.name,
+                input_types=udf._input_types,
+                return_type=udf._return_type,
+                original_return_type=original_return_type,
+            )
+        return snowpark_udf
+
+    snowpark_udf = get_snowpark_udf(udf_proto)
     udf_call_expr = snowpark_fn.call_udf(snowpark_udf.name, *snowpark_udf_args)
 
     # If the original return type was MapType or StructType but we converted it to VariantType,

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.'

snowflake/snowpark_connect/expression/map_unresolved_function.py

@@ -476,11 +476,8 @@ def map_unresolved_function(
             return TypedColumn(result, lambda: expected_types)
 
     match function_name:
-        case func_name if (
-            get_is_evaluating_sql() and func_name.lower() in session._udfs
-        ):
-            # TODO: In Spark, UDFs can override built-in functions in SQL,
-            # but not in DataFrame ops.
+        case func_name if func_name.lower() in session._udfs:
+            # In Spark, UDFs can override built-in functions
             udf = session._udfs[func_name.lower()]
             result_exp = snowpark_fn.call_udf(
                 udf.name,
@@ -714,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,
@@ -727,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:
@@ -2424,7 +2427,7 @@ def map_unresolved_function(
                 "try_to_date",
                 snowpark_fn.cast(
                     truncated_date,
-                    TimestampType(snowpark.types.TimestampTimeZone.NTZ),
+                    TimestampType(),
                 ),
                 snowpark_args[1],
             )
@@ -6479,6 +6482,18 @@ def map_unresolved_function(
                     if pattern_value is None:
                         return snowpark_fn.lit(None)
 
+                    # Optimization: treat escaped regex that resolves to a pure literal delimiter
+                    # - Single char: "\\."
+                    # - Multi char: e.g., "\\.505\\."
+                    if re.fullmatch(r"(?:\\.)+", pattern_value):
+                        literal_delim = re.sub(r"\\(.)", r"\1", pattern_value)
+                        return snowpark_fn.when(
+                            limit <= 0,
+                            snowpark_fn.split(
+                                str_, snowpark_fn.lit(literal_delim)
+                            ).cast(result_type),
+                        ).otherwise(native_split)
+
                     is_regexp = re.match(
                         ".*[\\[\\.\\]\\*\\?\\+\\^\\$\\{\\}\\|\\(\\)\\\\].*",
                         pattern_value,
@@ -8285,15 +8300,6 @@ def map_unresolved_function(
                 ),
             )
             result_type = BinaryType()
-        case udf_name if udf_name.lower() in session._udfs:
-            # TODO: In Spark, UDFs can override built-in functions in SQL,
-            # but not in DataFrame ops.
-            udf = session._udfs[udf_name.lower()]
-            result_exp = snowpark_fn.call_udf(
-                udf.name,
-                *(snowpark_fn.cast(arg, VariantType()) for arg in snowpark_args),
-            )
-            result_type = udf.return_type
         case udtf_name if udtf_name.lower() in session._udtfs:
             udtf, spark_col_names = session._udtfs[udtf_name.lower()]
             result_exp = snowpark_fn.call_table_function(
@@ -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
@@ -9623,13 +9629,14 @@ def _get_decimal_division_result_exp(
     snowpark_args: list[Column],
     spark_function_name: str,
 ) -> Column:
-    if isinstance(other_type, DecimalType) and overflow_detected:
-        if global_config.spark_sql_ansi_enabled:
-            raise ArithmeticException(
-                f'[NUMERIC_VALUE_OUT_OF_RANGE] {spark_function_name} cannot be represented as Decimal({result_type.precision}, {result_type.scale}). If necessary set "spark.sql.ansi.enabled" to "false" to bypass this error, and return NULL instead.'
-            )
-        else:
-            result_exp = snowpark_fn.lit(None).cast(result_type)
+    if (
+        isinstance(other_type, DecimalType)
+        and overflow_detected
+        and global_config.spark_sql_ansi_enabled
+    ):
+        raise ArithmeticException(
+            f'[NUMERIC_VALUE_OUT_OF_RANGE] {spark_function_name} cannot be represented as Decimal({result_type.precision}, {result_type.scale}). If necessary set "spark.sql.ansi.enabled" to "false" to bypass this error, and return NULL instead.'
+        )
     else:
         dividend = snowpark_args[0].cast(DoubleType())
         divisor = snowpark_args[1]
@@ -9910,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))"
@@ -9918,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:

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]] = []

snowflake/snowpark_connect/relation/map_map_partitions.py

@@ -8,28 +8,20 @@ from pyspark.sql.connect.proto.expressions_pb2 import CommonInlineUserDefinedFun
 import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
 from snowflake.snowpark.types import StructType
-from snowflake.snowpark_connect.config import global_config
 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.pandas_udtf_utils import create_pandas_udtf
-from snowflake.snowpark_connect.utils.udf_helper import (
-    SnowparkUDF,
-    process_udf_in_sproc,
-    require_creating_udf_in_sproc,
-    udf_check,
-)
-from snowflake.snowpark_connect.utils.udf_utils import (
-    ProcessCommonInlineUserDefinedFunction,
+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,
 )
-from snowflake.snowpark_connect.utils.udxf_import_utils import (
-    get_python_udxf_import_files,
-)
 
 
 def map_map_partitions(
@@ -41,17 +33,10 @@ def map_map_partitions(
     This is a simple wrapper around the `mapInPandas` method in Snowpark.
     """
     input_container = map_relation(rel.map_partitions.input)
-    input_df = input_container.dataframe
     udf_proto = rel.map_partitions.func
     udf_check(udf_proto)
 
-    # Check if this is mapInArrow (eval_type == 207)
-    eval_type = udf_proto.python_udf.eval_type
-
-    if eval_type == MAP_IN_ARROW_EVAL_TYPE:
-        return _map_in_arrow_with_pandas_udtf(input_container, udf_proto)
-    else:
-        return _map_partitions_with_udf(input_df, udf_proto)
+    return _map_with_pandas_udtf(input_container, udf_proto)
 
 
 def _call_udtf(
@@ -70,12 +55,17 @@ def _call_udtf(
 
     udtf_columns = input_df.columns + [snowpark_fn.col("_DUMMY_PARTITION_KEY")]
 
-    result_df_with_dummy = input_df_with_dummy.select(
-        snowpark_fn.call_table_function(udtf_name, *udtf_columns).over(
-            partition_by=[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.
@@ -89,7 +79,7 @@ def _call_udtf(
     )
 
 
-def _map_in_arrow_with_pandas_udtf(
+def _map_with_pandas_udtf(
     input_df_container: DataFrameContainer,
     udf_proto: CommonInlineUserDefinedFunction,
 ) -> snowpark.DataFrame:
@@ -99,55 +89,29 @@ def _map_in_arrow_with_pandas_udtf(
     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)
+    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:
-        map_in_arrow_udtf = create_pandas_udtf(
-            udf_proto, spark_column_names, input_schema, return_type
-        )
-        udtf_name = map_in_arrow_udtf.name
+        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)
-
-
-def _map_partitions_with_udf(
-    input_df: snowpark.DataFrame, udf_proto
-) -> snowpark.DataFrame:
-    """
-    Original UDF-based approach for non-mapInArrow map_partitions cases.
-    """
-    input_column_names = input_df.columns
-    kwargs = {
-        "common_inline_user_defined_function": udf_proto,
-        "input_types": [f.datatype for f in input_df.schema.fields],
-        "called_from": "map_map_partitions",
-        "udf_name": "spark_map_partitions_udf",
-        "input_column_names": input_column_names,
-        "replace": True,
-        "return_type": proto_to_snowpark_type(udf_proto.python_udf.output_type),
-        "udf_packages": global_config.get("snowpark.connect.udf.packages", ""),
-        "udf_imports": get_python_udxf_import_files(input_df.session),
-    }
-
-    if require_creating_udf_in_sproc(udf_proto):
-        snowpark_udf = process_udf_in_sproc(**kwargs)
-    else:
-        udf_processor = ProcessCommonInlineUserDefinedFunction(**kwargs)
-        udf = udf_processor.create_udf()
-        snowpark_udf = SnowparkUDF(
-            name=udf.name,
-            input_types=udf._input_types,
-            return_type=udf._return_type,
-            original_return_type=None,
-        )
-    udf_column_name = "UDF_OUTPUT"
-    snowpark_columns = [snowpark_fn.col(name) for name in input_df.columns]
-    result = input_df.select(snowpark_fn.call_udf(snowpark_udf.name, *snowpark_columns))
-    return DataFrameContainer.create_with_column_mapping(
-        dataframe=result,
-        spark_column_names=[udf_column_name],
-        snowpark_column_names=[udf_column_name],
-        snowpark_column_types=[snowpark_udf.return_type],
-    )