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

snowflake/snowpark_connect/expression/map_unresolved_star.py

@@ -34,7 +34,6 @@ def map_unresolved_star(
     column_mapping: ColumnNameMap,
     typer: ExpressionTyper,
 ) -> tuple[list[str], TypedColumn]:
-
     if exp.unresolved_star.HasField("unparsed_target"):
         unparsed_target = exp.unresolved_star.unparsed_target
         name_parts = split_fully_qualified_spark_name(unparsed_target)
@@ -103,7 +102,7 @@ def map_unresolved_star(
                 prefix_candidate_str = f"{prefix_candidate_str}.{name_parts[i]}"
             prefix_candidate = (
                 column_mapping.get_snowpark_column_name_from_spark_column_name(
-                    prefix_candidate_str, allow_non_exists=True, is_qualified=(i > 0)
+                    prefix_candidate_str, allow_non_exists=True
                 )
             )
             if prefix_candidate is None:
@@ -181,7 +180,7 @@ def map_unresolved_star_struct(
             prefix_candidate_str = f"{prefix_candidate_str}.{name_parts[i]}"
         prefix_candidate = (
             column_mapping.get_snowpark_column_name_from_spark_column_name(
-                prefix_candidate_str, allow_non_exists=True, is_qualified=(i > 0)
+                prefix_candidate_str, allow_non_exists=True
             )
         )
         if prefix_candidate is None:

snowflake/snowpark_connect/relation/map_extension.py

@@ -429,12 +429,18 @@ def map_aggregate(
             if groupings:
                 # Normal GROUP BY with explicit grouping columns
                 result = input_df.group_by(groupings)
-            else:
+            elif not is_group_by_all:
                 # No explicit GROUP BY - this is an aggregate over the entire table
                 # Use a dummy constant that will be excluded from the final result
                 result = input_df.with_column(
                     "__dummy_group__", snowpark_fn.lit(1)
                 ).group_by("__dummy_group__")
+            else:
+                # GROUP BY ALL with only one aggregate column
+                # Snowpark doesn't support GROUP BY ALL
+                # TODO: Change in future with Snowpark Supported arguments or API for GROUP BY ALL
+                result = input_df.group_by()
+
         case snowflake_proto.Aggregate.GROUP_TYPE_ROLLUP:
             result = input_df.rollup(groupings)
         case snowflake_proto.Aggregate.GROUP_TYPE_CUBE:

snowflake/snowpark_connect/relation/map_join.py

@@ -1,18 +1,13 @@
 #
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
-from collections import Counter
+
 from functools import reduce
 
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
-from pyspark.errors.exceptions.base import AnalysisException
 
 import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
-from snowflake.snowpark._internal.analyzer.analyzer_utils import (
-    quote_name_without_upper_casing,
-    unquote_if_quoted,
-)
 from snowflake.snowpark_connect.column_name_handler import JoinColumnNameMap
 from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.constants import COLUMN_METADATA_COLLISION_KEY
@@ -22,7 +17,6 @@ from snowflake.snowpark_connect.expression.map_expression import (
     map_single_column_expression,
 )
 from snowflake.snowpark_connect.expression.typer import JoinExpressionTyper
-from snowflake.snowpark_connect.hidden_column import HiddenColumn
 from snowflake.snowpark_connect.relation.map_relation import (
     NATURAL_JOIN_TYPE_BASE,
     map_relation,
@@ -30,6 +24,7 @@ from snowflake.snowpark_connect.relation.map_relation import (
 from snowflake.snowpark_connect.utils.context import (
     push_evaluating_join_condition,
     push_sql_scope,
+    set_plan_id_map,
     set_sql_plan_name,
 )
 from snowflake.snowpark_connect.utils.telemetry import (
@@ -38,9 +33,6 @@ from snowflake.snowpark_connect.utils.telemetry import (
 
 USING_COLUMN_NOT_FOUND_ERROR = "[UNRESOLVED_USING_COLUMN_FOR_JOIN] USING column `{0}` not found on the {1} side of the join. The {1}-side columns: {2}"
 
-DUPLICATED_JOIN_COL_LSUFFIX = "_left"
-DUPLICATED_JOIN_COL_RSUFFIX = "_right"
-
 
 def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
     left_container: DataFrameContainer = map_relation(rel.join.left)
@@ -82,13 +74,6 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
 
     # This handles case sensitivity for using_columns
     case_corrected_right_columns: list[str] = []
-    hidden_columns = set()
-    # Propagate the hidden columns from left/right inputs to the result in case of chained joins
-    if left_container.column_map.hidden_columns:
-        hidden_columns.update(left_container.column_map.hidden_columns)
-
-    if right_container.column_map.hidden_columns:
-        hidden_columns.update(right_container.column_map.hidden_columns)
 
     if rel.join.HasField("join_condition"):
         assert not using_columns
@@ -120,8 +105,8 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
             right=right_input,
             on=join_expression.col,
             how=join_type,
-            lsuffix=DUPLICATED_JOIN_COL_LSUFFIX,
-            rsuffix=DUPLICATED_JOIN_COL_RSUFFIX,
+            lsuffix="_left",
+            rsuffix="_right",
         )
     elif using_columns:
         if any(
@@ -171,24 +156,12 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
                 )
             )
 
-        using_columns_snowpark_names = (
-            left_container.column_map.get_snowpark_column_names_from_spark_column_names(
-                list(using_columns), return_first=True
-            )
-        )
-
-        using_columns_snowpark_types = [
-            left_container.dataframe.schema.fields[idx].datatype
-            for idx, col in enumerate(left_container.column_map.get_snowpark_columns())
-            if col in using_columns_snowpark_names
-        ]
-
         # Round trip the using columns through the column map to get the correct names
         # in order to support case sensitivity.
         # TODO: case_corrected_left_columns / case_corrected_right_columns may no longer be required as Snowpark dataframe preserves the column casing now.
-        case_corrected_left_columns = (
-            left_container.column_map.get_spark_column_names_from_snowpark_column_names(
-                using_columns_snowpark_names
+        case_corrected_left_columns = left_container.column_map.get_spark_column_names_from_snowpark_column_names(
+            left_container.column_map.get_snowpark_column_names_from_spark_column_names(
+                list(using_columns), return_first=True
             )
         )
         case_corrected_right_columns = right_container.column_map.get_spark_column_names_from_snowpark_column_names(
@@ -222,141 +195,28 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
                 (left == right for left, right in snowpark_using_columns),
             ),
             how=join_type,
-            rsuffix=DUPLICATED_JOIN_COL_RSUFFIX,
         )
-        # If we disambiguated the snowpark_using_columns during the join, we need to update 'snowpark_using_columns' to
-        # use the disambiguated names.
-        disambiguated_snowpark_using_columns = []
-
-        # Ignore disambiguation for LEFT SEMI JOIN and LEFT ANTI JOIN because they drop the right columns, so it'll never disambiguate.
-        if join_type in ["leftsemi", "leftanti"]:
-            disambiguated_snowpark_using_columns = snowpark_using_columns
-        else:
-            normalized_joined_columns = [
-                unquote_if_quoted(col) for col in joined_df.columns
-            ]
-            # snowpark_using_columns is a list of tuples of snowpark columns, joined_df.columns is a list of strings of column names
-            for (left, right) in snowpark_using_columns:
-                normalized_left_name = unquote_if_quoted(left.getName())
-                normalized_right_name = unquote_if_quoted(right.getName())
-
-                # are both left and right in joined_df? if not, it's been disambiguated
-                if (
-                    normalized_left_name in normalized_joined_columns
-                    and normalized_right_name in normalized_joined_columns
-                ):
-                    # we want to just add this
-                    disambiguated_snowpark_using_columns.append((left, right))
-                else:
-                    # we need to figure out the disambiguated names and add those - it only disambiguates if left == right
-                    disambiguated_left: snowpark.Column | None = None
-                    disambiguated_right: snowpark.Column | None = None
-
-                    for col in normalized_joined_columns:
-                        quoted_col = f'"{col}"'
-                        # get the column name and cross check it to see if it ends with the og name
-                        if col.endswith(normalized_left_name) and col.startswith("l_"):
-                            disambiguated_left = joined_df[quoted_col]
-                        elif col.endswith(normalized_right_name) and col.startswith(
-                            "r_"
-                        ):
-                            disambiguated_right = joined_df[quoted_col]
-
-                        # If we have both disambiguated columns, we can break out of the loop to save processing time
-                        if (
-                            disambiguated_left is not None
-                            and disambiguated_right is not None
-                        ):
-                            break
-                    if disambiguated_left is None or disambiguated_right is None:
-                        raise AnalysisException(
-                            f"Disambiguated columns not found for {normalized_left_name} and {normalized_right_name}."
-                        )
-                    disambiguated_snowpark_using_columns.append(
-                        (disambiguated_left, disambiguated_right)
-                    )
-
         # For outer joins, we need to preserve join keys from both sides using COALESCE
-        """
-        CHANGES:
-            - IF CASE
-                - Need to drop the using columns
-                - Need to create the hidden_columns DF with the using columns from right and left
-            - ELSE CASE
-                - Need to drop the right side using columns
-                - Need to create the hidden_columns DF with the using columns from right
-        """
         if join_type == "full_outer":
             coalesced_columns = []
-            for i, (left_col, _right_col) in enumerate(snowpark_using_columns):
+            columns_to_drop = []
+            for i, (left_col, right_col) in enumerate(snowpark_using_columns):
                 # Use the original user-specified column name to preserve case sensitivity
-                # Use the disambiguated columns for coalescing
-                disambiguated_left_col = disambiguated_snowpark_using_columns[i][0]
-                disambiguated_right_col = disambiguated_snowpark_using_columns[i][1]
-
-                coalesced_col = snowpark_fn.coalesce(
-                    disambiguated_left_col, disambiguated_right_col
-                ).alias(left_col.get_name())
-                coalesced_columns.append(coalesced_col)
-
-                # Create HiddenColumn objects for each hidden column
-                hidden_left = HiddenColumn(
-                    hidden_snowpark_name=disambiguated_left_col.getName(),
-                    spark_name=case_corrected_left_columns[i],
-                    visible_snowpark_name=left_col.get_name(),
-                    qualifiers=left_container.column_map.get_qualifier_for_spark_column(
-                        case_corrected_left_columns[i]
-                    ),
-                    original_position=left_container.column_map.get_spark_columns().index(
-                        case_corrected_left_columns[i]
-                    ),
-                )
-
-                hidden_right = HiddenColumn(
-                    hidden_snowpark_name=disambiguated_right_col.getName(),
-                    spark_name=case_corrected_right_columns[i],
-                    visible_snowpark_name=left_col.get_name(),
-                    qualifiers=right_container.column_map.get_qualifier_for_spark_column(
-                        case_corrected_right_columns[i]
-                    ),
-                    original_position=right_container.column_map.get_spark_columns().index(
-                        case_corrected_right_columns[i]
-                    ),
-                )
-                hidden_columns.update(
-                    [
-                        hidden_left,
-                        hidden_right,
-                    ]
+                original_column_name = rel.join.using_columns[i]
+                coalesced_col = snowpark_fn.coalesce(left_col, right_col).alias(
+                    original_column_name
                 )
+                coalesced_columns.append(coalesced_col)
+                columns_to_drop.extend([left_col, right_col])
 
-            # All non-hidden columns (not including the coalesced columns)
             other_columns = [
                 snowpark_fn.col(col_name)
                 for col_name in joined_df.columns
-                if col_name not in [col.hidden_snowpark_name for col in hidden_columns]
+                if col_name not in [col.getName() for col in columns_to_drop]
             ]
             result = joined_df.select(coalesced_columns + other_columns)
-
         else:
             result = joined_df.drop(*(right for _, right in snowpark_using_columns))
-            # We never run into the disambiguation case unless it's a full outer join.
-            for i, (left_col, right_col) in enumerate(
-                disambiguated_snowpark_using_columns
-            ):
-                # Only right side columns are hidden
-                hidden_col = HiddenColumn(
-                    hidden_snowpark_name=right_col.getName(),
-                    spark_name=case_corrected_right_columns[i],
-                    visible_snowpark_name=left_col.getName(),
-                    qualifiers=right_container.column_map.get_qualifier_for_spark_column(
-                        case_corrected_right_columns[i]
-                    ),
-                    original_position=right_container.column_map.get_spark_columns().index(
-                        case_corrected_right_columns[i]
-                    ),
-                )
-                hidden_columns.add(hidden_col)
     else:
         if join_type != "cross" and not global_config.spark_sql_crossJoin_enabled:
             raise SparkException.implicit_cartesian_product("inner")
@@ -370,110 +230,35 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
         # - LEFT SEMI JOIN: Returns left rows that have matches in right table (no right columns)
         # - LEFT ANTI JOIN: Returns left rows that have NO matches in right table (no right columns)
         # Both preserve only the columns from the left DataFrame without adding any columns from the right.
-        spark_cols_after_join = left_container.column_map.get_spark_columns()
-        snowpark_cols_after_join = left_container.column_map.get_snowpark_columns()
-        snowpark_col_types = [
-            f.datatype for f in left_container.dataframe.schema.fields
-        ]
+        spark_cols_after_join: list[str] = left_container.column_map.get_spark_columns()
         qualifiers = left_container.column_map.get_qualifiers()
-    elif join_type == "full_outer" and using_columns:
-        # We want the coalesced columns to be first, followed by all the left and right columns (excluding using columns)
-        spark_cols_after_join: list[str] = []
-        snowpark_cols_after_join: list[str] = []
-        snowpark_col_types: list[str] = []
-
-        left_container_snowpark_columns = (
-            left_container.column_map.get_snowpark_columns()
-        )
-        right_container_snowpark_columns = (
-            right_container.column_map.get_snowpark_columns()
-        )
-
-        qualifiers = []
-        for i in range(len(case_corrected_left_columns)):
-            spark_cols_after_join.append(case_corrected_left_columns[i])
-            snowpark_cols_after_join.append(using_columns_snowpark_names[i])
-            snowpark_col_types.append(using_columns_snowpark_types[i])
-            qualifiers.append([])
-
-        # Handle adding left and right columns, excluding the using columns
-        for i, spark_col in enumerate(left_container.column_map.get_spark_columns()):
-            if (
-                spark_col not in case_corrected_left_columns
-                or spark_col in left_container.column_map.get_spark_columns()[:i]
-            ):
-                spark_cols_after_join.append(spark_col)
-                snowpark_cols_after_join.append(left_container_snowpark_columns[i])
-                qualifiers.append(
-                    left_container.column_map.get_qualifier_for_spark_column(spark_col)
-                )
-
-                snowpark_col_types.append(
-                    left_container.dataframe.schema.fields[i].datatype
-                )
-
-        for i, spark_col in enumerate(right_container.column_map.get_spark_columns()):
-            if (
-                spark_col not in case_corrected_right_columns
-                or spark_col in right_container.column_map.get_spark_columns()[:i]
-            ):
-                spark_cols_after_join.append(spark_col)
-                snowpark_cols_after_join.append(right_container_snowpark_columns[i])
-                qualifiers.append(
-                    right_container.column_map.get_qualifier_for_spark_column(spark_col)
-                )
-
-                snowpark_col_types.append(
-                    right_container.dataframe.schema.fields[i].datatype
-                )
-
     else:
-        spark_cols_after_join = left_container.column_map.get_spark_columns()
-        snowpark_cols_after_join = left_container.column_map.get_snowpark_columns()
-        snowpark_col_types = [
-            f.datatype for f in left_container.dataframe.schema.fields
+        # Add Spark columns and plan_ids from left DF
+        spark_cols_after_join: list[str] = list(
+            left_container.column_map.get_spark_columns()
+        ) + [
+            spark_col
+            for i, spark_col in enumerate(
+                right_container.column_map.get_spark_columns()
+            )
+            if spark_col not in case_corrected_right_columns
+            or spark_col
+            in right_container.column_map.get_spark_columns()[
+                :i
+            ]  # this is to make sure we only remove the column once
         ]
 
-        qualifiers = left_container.column_map.get_qualifiers()
-
-        right_df_snowpark_columns = right_container.column_map.get_snowpark_columns()
-
-        for i, spark_col in enumerate(right_container.column_map.get_spark_columns()):
-            if (
-                spark_col not in case_corrected_right_columns
-                or spark_col in right_container.column_map.get_spark_columns()[:i]
-            ):
-                spark_cols_after_join.append(spark_col)
-                snowpark_cols_after_join.append(right_df_snowpark_columns[i])
-                snowpark_col_types.append(
-                    right_container.dataframe.schema.fields[i].datatype
-                )
-
-                qualifiers.append(
-                    right_container.column_map.get_qualifier_for_spark_column(spark_col)
-                )
-
-    snowpark_cols_after_join_deduplicated = []
-    snowpark_cols_after_join_counter = Counter(snowpark_cols_after_join)
-    seen_duplicated_columns = set()
-
-    for col in snowpark_cols_after_join:
-        if snowpark_cols_after_join_counter[col] == 2:
-            # This means that the same column exists twice in the joined df, likely due to a self-join and
-            # we need to lsuffix and rsuffix to the names of both columns, similar to what Snowpark did under the hood.
-
-            suffix = (
-                DUPLICATED_JOIN_COL_RSUFFIX
-                if col in seen_duplicated_columns
-                else DUPLICATED_JOIN_COL_LSUFFIX
+        qualifiers = list(left_container.column_map.get_qualifiers()) + [
+            right_container.column_map.get_qualifier_for_spark_column(spark_col)
+            for i, spark_col in enumerate(
+                right_container.column_map.get_spark_columns()
             )
-            unquoted_col = unquote_if_quoted(col)
-            quoted = quote_name_without_upper_casing(unquoted_col + suffix)
-            snowpark_cols_after_join_deduplicated.append(quoted)
-
-            seen_duplicated_columns.add(col)
-        else:
-            snowpark_cols_after_join_deduplicated.append(col)
+            if spark_col not in case_corrected_right_columns
+            or spark_col
+            in right_container.column_map.get_spark_columns()[
+                :i
+            ]  # this is to make sure we only remove the column once]
+        ]
 
     column_metadata = {}
     if left_container.column_map.column_metadata:
@@ -502,13 +287,33 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
     result_container = DataFrameContainer.create_with_column_mapping(
         dataframe=result,
         spark_column_names=spark_cols_after_join,
-        snowpark_column_names=snowpark_cols_after_join_deduplicated,
+        snowpark_column_names=result.columns,
         column_metadata=column_metadata,
         column_qualifiers=qualifiers,
-        hidden_columns=hidden_columns,
-        snowpark_column_types=snowpark_col_types,
     )
 
+    # Fix for USING join column references with different plan IDs
+    # After a USING join, references to the right dataframe's columns should resolve
+    # to the result dataframe that contains the merged columns
+    if (
+        using_columns
+        and rel.join.right.HasField("common")
+        and rel.join.right.common.HasField("plan_id")
+    ):
+        right_plan_id = rel.join.right.common.plan_id
+        set_plan_id_map(right_plan_id, result_container)
+
+    # For FULL OUTER joins, we also need to map the left dataframe's plan_id
+    # since both columns are replaced with a coalesced column
+    if (
+        using_columns
+        and join_type == "full_outer"
+        and rel.join.left.HasField("common")
+        and rel.join.left.common.HasField("plan_id")
+    ):
+        left_plan_id = rel.join.left.common.plan_id
+        set_plan_id_map(left_plan_id, result_container)
+
     if rel.join.using_columns:
         # When join 'using_columns', the 'join columns' should go first in result DF.
         idxs_to_shift = [
@@ -540,7 +345,6 @@ def map_join(rel: relation_proto.Relation) -> DataFrameContainer:
             cached_schema_getter=lambda: snowpark.types.StructType(
                 reorder(original_df.schema.fields)
             ),
-            hidden_columns=hidden_columns,
         )
 
     return result_container

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,18 +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)
-    if (
-        udf_proto.WhichOneof("function") == "python_udf"
-        and udf_proto.python_udf.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(
@@ -71,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.
@@ -90,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:
@@ -100,59 +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
-            if udf_proto.WhichOneof("function") == "python_udf"
-            else udf_proto.scalar_scala_udf.outputType
-        ),
-        "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],
-    )

snowflake/snowpark_connect/relation/map_relation.py

@@ -15,6 +15,8 @@ from snowflake.snowpark_connect.utils.cache import (
 from snowflake.snowpark_connect.utils.context import (
     get_plan_id_map,
     get_session_id,
+    not_resolving_fun_args,
+    push_map_partitions,
     push_operation_scope,
     set_is_aggregate_function,
     set_plan_id_map,
@@ -149,7 +151,10 @@ def map_relation(
             case "drop_na":
                 result = map_row_ops.map_dropna(rel)
             case "extension":
-                result = map_extension.map_extension(rel)
+                # Extensions can be passed as function args, and we need to reset the context here.
+                # Matters only for resolving alias expressions in the extensions rel.
+                with not_resolving_fun_args():
+                    result = map_extension.map_extension(rel)
             case "fill_na":
                 result = map_row_ops.map_fillna(rel)
             case "filter":
@@ -180,7 +185,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":