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

snowflake/snowpark_connect/column_name_handler.py

@@ -20,7 +20,11 @@ from snowflake.snowpark._internal.analyzer.analyzer_utils import (
 from snowflake.snowpark._internal.utils import quote_name
 from snowflake.snowpark.types import StructType
 from snowflake.snowpark_connect.config import global_config
+from snowflake.snowpark_connect.hidden_column import HiddenColumn
 from snowflake.snowpark_connect.utils.context import get_current_operation_scope
+from snowflake.snowpark_connect.utils.identifiers import (
+    split_fully_qualified_spark_name,
+)
 
 ALREADY_QUOTED = re.compile('^(".+")$', re.DOTALL)
 
@@ -84,6 +88,8 @@ class ColumnNames:
     spark_name: str
     snowpark_name: str
     qualifiers: list[str]
+    catalog_info: str | None = None  # Catalog from fully qualified name
+    database_info: str | None = None  # Database from fully qualified name
 
 
 def get_list_of_spark_names_for_column(column_names: ColumnNames) -> list[str]:
@@ -118,6 +124,7 @@ class ColumnNameMap:
         ] = lambda: global_config.spark_sql_caseSensitive,
         column_metadata: dict | None = None,
         column_qualifiers: list[list[str]] | None = None,
+        hidden_columns: set[HiddenColumn] | None = None,
         parent_column_name_map: ColumnNameMap | None = None,
     ) -> None:
         """
@@ -128,6 +135,7 @@ class ColumnNameMap:
         The key is the original Spark column name, and the value is the metadata.
         example: Dict('age', {'foo': 'bar'})
         column_qualifiers: Optional qualifiers for the columns, used to handle table aliases or DataFrame aliases.
+        hidden_columns: Optional set of HiddenColumn objects.
         parent_column_name_map: parent ColumnNameMap
         """
         self.columns: list[ColumnNames] = []
@@ -136,6 +144,7 @@ class ColumnNameMap:
         self.snowpark_to_col = defaultdict(list)
         self.is_case_sensitive = is_case_sensitive
         self.column_metadata = column_metadata
+        self.hidden_columns = hidden_columns
 
         # Rename chain dictionary to track column renaming history
         self.rename_chains: dict[str, str] = {}  # old_name -> new_name mapping
@@ -152,10 +161,27 @@ class ColumnNameMap:
         )
 
         for i in range(len(spark_column_names)):
+            # Extract catalog/database info from fully qualified spark column names
+            spark_name = spark_column_names[i]
+            catalog_info = None
+            database_info = None
+
+            # Parse fully qualified names to extract catalog.database.table.column format
+            name_parts = split_fully_qualified_spark_name(spark_name)
+            if len(name_parts) >= 4:
+                # Format: catalog.database.table.column[.field...]
+                # Only extract if this looks like a valid SQL table reference
+                # Catalog.database.column (3 parts) is invalid - missing table
+                # Only catalog.database.table.column (4+ parts) is potentially valid
+                catalog_info = name_parts[0]
+                database_info = name_parts[1]
+
             c = ColumnNames(
-                spark_name=spark_column_names[i],
+                spark_name=spark_name,
                 snowpark_name=snowpark_column_names[i],
                 qualifiers=column_qualifiers[i] if column_qualifiers else [],
+                catalog_info=catalog_info,
+                database_info=database_info,
             )
             self.columns.append(c)
 
@@ -312,6 +338,8 @@ class ColumnNameMap:
         *,
         allow_non_exists: bool = False,
         return_first: bool = False,
+        is_qualified: bool = False,
+        source_qualifiers: list[str] | None = None,
     ) -> str | None:
         assert isinstance(spark_column_name, str)
         resolved_name = (
@@ -319,9 +347,37 @@ class ColumnNameMap:
             if self.rename_chains
             else spark_column_name
         )
-        snowpark_names = self.get_snowpark_column_names_from_spark_column_names(
-            [resolved_name], return_first
-        )
+
+        # We need to check hidden columns first. We want to avoid the code path
+        # within get_snowpark_column_names_from_spark_column_names that checks the parent ColumnNameMap.
+        # This is because that will return the name of the using column that's been dropped from the result
+        # dataframe. We want to fetch and resolve the hidden column to its visible using column name instead.
+        # Even if this is an unqualified reference or one to the visible column, it will resolve correctly to
+        # the visible name anyway.
+        snowpark_names = []
+        # Only check hidden columns for qualified references with source qualifiers
+        if is_qualified and source_qualifiers is not None and self.hidden_columns:
+            column_name = spark_column_name
+
+            # Check each hidden column for column name AND qualifier match
+            for hidden_col in self.hidden_columns:
+                if (
+                    hidden_col.spark_name == column_name
+                    and hidden_col.qualifiers == source_qualifiers
+                ):
+                    if not global_config.spark_sql_caseSensitive:
+                        if hidden_col.spark_name.upper() == column_name.upper() and [
+                            q.upper() for q in hidden_col.qualifiers
+                        ] == [q.upper() for q in source_qualifiers]:
+                            snowpark_names.append(hidden_col.visible_snowpark_name)
+                    else:
+                        snowpark_names.append(hidden_col.visible_snowpark_name)
+
+        # If not found in hidden columns, proceed with normal lookup
+        if not snowpark_names:
+            snowpark_names = self.get_snowpark_column_names_from_spark_column_names(
+                [resolved_name], return_first
+            )
 
         snowpark_names_len = len(snowpark_names)
         if snowpark_names_len > 1:
@@ -408,6 +464,27 @@ class ColumnNameMap:
                 snowpark_columns.append(c.snowpark_name)
                 qualifiers.append(c.qualifiers)
 
+        # Note: The following code is commented out because there is a bug with handling duplicate columns in
+        # qualified select *'s. This needs to be revisited once a solution for that is found.
+        # TODO: https://snowflakecomputing.atlassian.net/browse/SNOW-2265240
+
+        # # Handles fetching/resolving the hidden columns if they also match the qualifiers
+        # # This method is only ever called for qualified references, so we need to check hidden columns as well.
+        # if self.hidden_columns:
+        #     for hidden_col in self.hidden_columns:
+        #         col_qualifiers = (
+        #             [q.upper() for q in hidden_col.qualifiers]
+        #             if not self.is_case_sensitive()
+        #             else hidden_col.qualifiers
+        #         )
+        #         if len(col_qualifiers) < len(qualifiers_input):
+        #             continue
+        #         if col_qualifiers[-len(qualifiers_input) :] == qualifiers_input:
+        #             # This hidden column matches! Add it to the results
+        #             spark_columns.append(hidden_col.spark_name)
+        #             snowpark_columns.append(hidden_col.visible_snowpark_name)
+        #             qualifiers.append(hidden_col.qualifiers)
+
         return spark_columns, snowpark_columns, qualifiers
 
     def get_snowpark_columns(self) -> list[str]:
@@ -539,6 +616,35 @@ class ColumnNameMap:
         else:
             return spark_name.upper()
 
+    def is_hidden_column_reference(
+        self, spark_column_name: str, source_qualifiers: list[str] | None = None
+    ) -> bool:
+        """
+        Check if a column reference would be resolved through hidden columns.
+        """
+        if not self.hidden_columns or source_qualifiers is None:
+            return False
+
+        # For qualified references with source_qualifiers
+        column_name = (
+            spark_column_name  # When has_plan_id=True, this is just the column name
+        )
+
+        for hidden_col in self.hidden_columns:
+            if (
+                hidden_col.spark_name == column_name
+                and hidden_col.qualifiers == source_qualifiers
+            ):
+                if not global_config.spark_sql_caseSensitive:
+                    if hidden_col.spark_name.upper() == column_name.upper() and [
+                        q.upper() for q in hidden_col.qualifiers
+                    ] == [q.upper() for q in source_qualifiers]:
+                        return True
+                else:
+                    return True
+
+        return False
+
 
 class JoinColumnNameMap(ColumnNameMap):
     def __init__(
@@ -548,6 +654,9 @@ class JoinColumnNameMap(ColumnNameMap):
     ) -> None:
         self.left_column_mapping: ColumnNameMap = left_colmap
         self.right_column_mapping: ColumnNameMap = right_colmap
+        # Ensure attributes expected by base-class helpers exist to avoid AttributeError
+        # when generic code paths (e.g., hidden column checks) touch them.
+        self.hidden_columns: set[HiddenColumn] | None = None
 
     def get_snowpark_column_name_from_spark_column_name(
         self,
@@ -555,6 +664,9 @@ class JoinColumnNameMap(ColumnNameMap):
         *,
         allow_non_exists: bool = False,
         return_first: bool = False,
+        # JoinColumnNameMap will never be called with using columns, so these parameters are not used.
+        is_qualified: bool = False,
+        source_qualifiers: list[str] | None = None,
     ) -> str | None:
         snowpark_column_name_in_left = (
             self.left_column_mapping.get_snowpark_column_name_from_spark_column_name(

snowflake/snowpark_connect/config.py

@@ -19,6 +19,7 @@ from snowflake.snowpark._internal.analyzer.analyzer_utils import (
     quote_name_without_upper_casing,
 )
 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.session import get_or_create_snowpark_session
@@ -593,3 +594,15 @@ def parse_imports(session: snowpark.Session, imports: str | None) -> None:
 
     for udf_import in imports.strip("[] ").split(","):
         session.add_import(udf_import)
+
+
+def get_timestamp_type():
+    match global_config["spark.sql.timestampType"]:
+        case "TIMESTAMP_LTZ":
+            timestamp_type = TimestampType(TimestampTimeZone.LTZ)
+        case "TIMESTAMP_NTZ":
+            timestamp_type = TimestampType(TimestampTimeZone.NTZ)
+        case _:
+            # shouldn't happen since `spark.sql.timestampType` is always defined, and `spark.conf.unset` sets it to default (TIMESTAMP_LTZ)
+            timestamp_type = TimestampType(TimestampTimeZone.LTZ)
+    return timestamp_type

snowflake/snowpark_connect/constants.py

@@ -13,35 +13,6 @@ STRUCTURED_TYPES_ENABLED = True
 # UDF evaluation types
 MAP_IN_ARROW_EVAL_TYPE = 207  # eval_type for mapInArrow operations
 
-SPARK_TZ_ABBREVIATIONS_OVERRIDES = {
-    "ACT": "Australia/Darwin",
-    "AET": "Australia/Sydney",
-    "AGT": "America/Argentina/Buenos_Aires",
-    "ART": "Africa/Cairo",
-    "AST": "America/Anchorage",
-    "BET": "America/Sao_Paulo",
-    "BST": "Asia/Dhaka",
-    "CAT": "Africa/Harare",
-    "CNT": "America/St_Johns",
-    "CST": "America/Chicago",
-    "CTT": "Asia/Shanghai",
-    "EAT": "Africa/Addis_Ababa",
-    "ECT": "Europe/Paris",
-    "IET": "America/Indiana/Indianapolis",
-    "IST": "Asia/Kolkata",
-    "JST": "Asia/Tokyo",
-    "MIT": "Pacific/Apia",
-    "NET": "Asia/Yerevan",
-    "NST": "Pacific/Auckland",
-    "PLT": "Asia/Karachi",
-    "PNT": "America/Phoenix",
-    "PRT": "America/Puerto_Rico",
-    "PST": "America/Los_Angeles",
-    "SST": "Pacific/Guadalcanal",
-    "VST": "Asia/Ho_Chi_Minh",
-}
-
-
 COLUMN_METADATA_COLLISION_KEY = "{expr_id}_{key}"
 
 DUPLICATE_KEY_FOUND_ERROR_TEMPLATE = "Duplicate key found: {key}. You can set spark.sql.mapKeyDedupPolicy to LAST_WIN to deduplicate map keys with last wins policy."

snowflake/snowpark_connect/dataframe_container.py

@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Callable
 
 from snowflake import snowpark
 from snowflake.snowpark.types import StructField, StructType
+from snowflake.snowpark_connect.hidden_column import HiddenColumn
 
 if TYPE_CHECKING:
     from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
@@ -57,6 +58,7 @@ class DataFrameContainer:
         column_metadata: dict | None = None,
         column_qualifiers: list[list[str]] | None = None,
         parent_column_name_map: ColumnNameMap | None = None,
+        hidden_columns: set[HiddenColumn] | None = None,
         table_name: str | None = None,
         alias: str | None = None,
         cached_schema_getter: Callable[[], StructType] | None = None,
@@ -72,6 +74,7 @@ class DataFrameContainer:
             column_metadata: Optional metadata dictionary
             column_qualifiers: Optional column qualifiers
             parent_column_name_map: Optional parent column name map
+            hidden_columns: Optional list of hidden column names
             table_name: Optional table name
             alias: Optional alias
             cached_schema_getter: Optional function to get cached schema
@@ -93,6 +96,7 @@ class DataFrameContainer:
             column_metadata,
             column_qualifiers,
             parent_column_name_map,
+            hidden_columns,
         )
 
         # Determine the schema getter to use
@@ -206,6 +210,7 @@ class DataFrameContainer:
         column_metadata: dict | None = None,
         column_qualifiers: list[list[str]] | None = None,
         parent_column_name_map: ColumnNameMap | None = None,
+        hidden_columns: set[HiddenColumn] | None = None,
     ) -> ColumnNameMap:
         """Create a ColumnNameMap with the provided configuration."""
         from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
@@ -216,6 +221,7 @@ class DataFrameContainer:
             column_metadata=column_metadata,
             column_qualifiers=column_qualifiers,
             parent_column_name_map=parent_column_name_map,
+            hidden_columns=hidden_columns,
         )
 
     @staticmethod

snowflake/snowpark_connect/execute_plan/map_execution_command.py

@@ -1,11 +1,15 @@
 #
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
+import re
+import uuid
 from collections import Counter
 
 import pyspark.sql.connect.proto.base_pb2 as proto_base
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
+from snowflake.snowpark import DataFrame, Session
+from snowflake.snowpark.exceptions import SnowparkSQLException
 from snowflake.snowpark_connect.column_name_handler import ColumnNames
 from snowflake.snowpark_connect.config import global_config, sessions_config
 from snowflake.snowpark_connect.constants import SERVER_SIDE_SESSION_ID
@@ -25,6 +29,10 @@ from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
 )
 
+_INTERNAL_VIEW_PREFIX = "__SC_RENAMED_V_"
+
+_CREATE_VIEW_PATTERN = re.compile(r"create\s+or\s+replace\s+view", re.IGNORECASE)
+
 
 def _create_column_rename_map(
     columns: list[ColumnNames], rename_duplicated: bool
@@ -126,7 +134,14 @@ def map_execution_command(
             ]
 
             if req.replace:
-                input_df.create_or_replace_temp_view(view_name)
+                try:
+                    input_df.create_or_replace_temp_view(view_name)
+                except SnowparkSQLException as exc:
+                    if _is_error_caused_by_view_referencing_itself(exc):
+                        # This error is caused by statement with self reference like `CREATE VIEW A AS SELECT X FROM A`.
+                        _create_chained_view(input_df, view_name)
+                    else:
+                        raise
             else:
                 input_df.create_temp_view(view_name)
         case "write_stream_operation_start":
@@ -192,3 +207,43 @@ def map_execution_command(
             raise SnowparkConnectNotImplementedError(
                 f"Command type {other} not implemented"
             )
+
+
+def _generate_random_builtin_view_name() -> str:
+    return _INTERNAL_VIEW_PREFIX + str(uuid.uuid4()).replace("-", "")
+
+
+def _is_error_caused_by_view_referencing_itself(exc: Exception) -> bool:
+    return "view definition refers to view being defined" in str(exc).lower()
+
+
+def _create_chained_view(input_df: DataFrame, view_name: str) -> None:
+    """
+    In order to create a view, which references itself, Spark would here take the previous
+    definition of A and paste it in place of `FROM A`. Snowflake would fail in such case, so
+    as a workaround, we create a chain of internal views instead. This function:
+    1. Renames previous definition of A to some internal name (instead of deleting).
+    2. Adjusts the DDL of a new statement to reference the name of a renmaed internal view, instead of itself.
+    """
+
+    session = Session.get_active_session()
+
+    view_name = ".".join(view_name)
+
+    tmp_name = _generate_random_builtin_view_name()
+    old_name_replacement = _generate_random_builtin_view_name()
+
+    input_df.create_or_replace_temp_view(tmp_name)
+
+    session.sql(f"ALTER VIEW {view_name} RENAME TO {old_name_replacement}").collect()
+
+    ddl: str = session.sql(f"SELECT GET_DDL('VIEW', '{tmp_name}')").collect()[0][0]
+
+    ddl = ddl.replace(view_name, old_name_replacement)
+
+    # GET_DDL result doesn't contain `TEMPORARY`, it's likely a bug.
+    ddl = _CREATE_VIEW_PATTERN.sub("create or replace temp view", ddl)
+
+    session.sql(ddl).collect()
+
+    session.sql(f"ALTER VIEW {tmp_name} RENAME TO {view_name}").collect()

snowflake/snowpark_connect/expression/literal.py

@@ -68,13 +68,24 @@ def get_literal_field_and_name(literal: expressions_proto.Expression.Literal):
                 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')}'"
+            def _format_timestamp(dt) -> str:
+                without_micros = f"{dt.year:04d}-{dt.month:02d}-{dt.day:02d} {dt.hour:02d}:{dt.minute:02d}:{dt.second:02d}"
+                if dt.microsecond == 0:
+                    return without_micros
+                else:
+                    base_format = f"{without_micros}.{dt.microsecond:06d}"
+                    return base_format.rstrip("0").rstrip(".")
+
+            return lit_dt, f"{t.upper()} '{_format_timestamp(tz_dt)}'"
         case "day_time_interval":
             # TODO(SNOW-1920942): Snowflake SQL is missing an "interval" type.
             timedelta = datetime.timedelta(
                 seconds=literal.day_time_interval / 1_000_000
             )
-            str_value = f"INTERVAL '{literal.day_time_interval / 1_000_000} SECONDS'"
+            interval_seconds = f"{(literal.day_time_interval / 1_000_000):.6f}".rstrip(
+                "0"
+            ).rstrip(".")
+            str_value = f"INTERVAL '{interval_seconds}' SECOND"
             return timedelta, str_value
         case "binary":
             return literal.binary, str(literal.binary)

snowflake/snowpark_connect/expression/map_cast.py

@@ -82,8 +82,10 @@ def map_cast(
     match exp.cast.WhichOneof("cast_to_type"):
         case "type":
             to_type = proto_to_snowpark_type(exp.cast.type)
+            to_type_str = to_type.simpleString().upper()
         case "type_str":
             to_type = map_type_string_to_snowpark_type(exp.cast.type_str)
+            to_type_str = exp.cast.type_str.upper()
         case _:
             raise ValueError("No type to cast to")
 
@@ -102,7 +104,7 @@ def map_cast(
         case _ if to_type.typeName().upper() in ("STRUCT", "ARRAY"):
             col_name = new_name
         case _ if get_is_evaluating_sql():
-            col_name = f"CAST({new_name} AS {to_type.simpleString().upper()})"
+            col_name = f"CAST({new_name} AS {to_type_str})"
         case _:
             col_name = new_name
 
@@ -121,17 +123,13 @@ def map_cast(
             sanity_check(to_type, new_name, from_type, from_type_cast)
 
     col = typed_column.col
-    # On TCM, sometimes these are StringType(x) and TimestampType(tz=x)
+    # On TCM, sometimes these are StringType(x)
     # This normalizes them for the cast.
     if isinstance(from_type, StringType):
         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()
-
     match (from_type, to_type):
         case (_, _) if (from_type == to_type):
             result_exp = col
@@ -196,9 +194,8 @@ def map_cast(
             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
+            result_exp = col.cast(to_type)
         case (_, TimestampType()) if isinstance(from_type, _NumericType):
             microseconds = col * snowpark_fn.lit(1000000)
             result_exp = snowpark_fn.when(

snowflake/snowpark_connect/expression/map_sql_expression.py

@@ -15,6 +15,7 @@ from pyspark.errors.exceptions.base import AnalysisException
 import snowflake.snowpark_connect.proto.snowflake_expression_ext_pb2 as snowflake_proto
 from snowflake import snowpark
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
+from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.typed_column import TypedColumn
 from snowflake.snowpark_connect.utils.context import (
     get_sql_named_arg,
@@ -34,11 +35,32 @@ DECIMAL_RE = re.compile(r"decimal\((\d+), *(\d+)\)")
 _window_specs = ContextVar[dict[str, any]]("_window_specs", default={})
 
 
-@cache
 def sql_parser():
+    """
+    SparkSqlParser needs spark.config. Here we are setting required config keys for SAS.
+    We can't do it in set.config request handler, because SparkSqlParser initialized without active SparkSession (our case, todo: fix it)
+    uses ThreadLocal to access conf, and conf.set can be handled by other thread.
+    https://github.com/apache/spark/blob/ea53ea71461508801586b1e5677aa6011df7cd95/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala#L143
+    """
+
+    ts_type = global_config.spark_sql_timestampType
+
+    if ts_type is not None:
+        _get_sql_conf().get().setConfString("spark.sql.timestampType", str(ts_type))
+
+    return _get_sql_parser()
+
+
+@cache
+def _get_sql_parser():
     return jpype.JClass("org.apache.spark.sql.execution.SparkSqlParser")()
 
 
+@cache
+def _get_sql_conf():
+    return jpype.JClass("org.apache.spark.sql.internal.SQLConf")
+
+
 @cache
 def _as_java_list():
     return jpype.JClass("scala.collection.JavaConverters").seqAsJavaList

snowflake/snowpark_connect/expression/map_udf.py

@@ -53,9 +53,16 @@ def register_udf(
     udf_proto: expressions_proto.CommonInlineUserDefinedFunction,
 ) -> SnowparkUDF:
     udf_check(udf_proto)
-    processed_return_type, original_return_type = process_udf_return_type(
-        udf_proto.python_udf.output_type
-    )
+    match udf_proto.WhichOneof("function"):
+        case "python_udf":
+            output_type = udf_proto.python_udf.output_type
+        case "scalar_scala_udf":
+            output_type = udf_proto.scalar_scala_udf.outputType
+        case _:
+            raise ValueError(
+                f"Unsupported UDF type: {udf_proto.WhichOneof('function')}"
+            )
+    processed_return_type, original_return_type = process_udf_return_type(output_type)
     session = get_or_create_snowpark_session()
     kwargs = {
         "common_inline_user_defined_function": udf_proto,
@@ -79,7 +86,10 @@ def register_udf(
         )
         # 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
+        if udf_processor._function_type == "scalar_scala_udf":
+            session._udfs[udf.name] = udf
+        else:
+            session._udfs[udf_proto.function_name.lower()] = udf
         return udf
 
 
@@ -94,9 +104,16 @@ def map_common_inline_user_defined_udf(
         udf_proto, column_mapping, typer
     )
     input_types = gen_input_types(snowpark_udf_args, typer)
-    processed_return_type, original_return_type = process_udf_return_type(
-        udf_proto.python_udf.output_type
-    )
+    match udf_proto.WhichOneof("function"):
+        case "python_udf":
+            processed_return_type, original_return_type = process_udf_return_type(
+                udf_proto.python_udf.output_type
+            )
+        case "scalar_scala_udf":
+            processed_return_type, original_return_type = process_udf_return_type(
+                udf_proto.scalar_scala_udf.outputType
+            )
+
     session = get_or_create_snowpark_session()
     kwargs = {
         "common_inline_user_defined_function": udf_proto,
@@ -118,7 +135,8 @@ def map_common_inline_user_defined_udf(
             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
     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,