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

snowflake/snowpark_connect/utils/udf_helper.py

@@ -14,14 +14,20 @@ from pyspark.errors.exceptions.base import AnalysisException
 import snowflake.snowpark.functions as snowpark_fn
 import snowflake.snowpark_connect.tcm as tcm
 import snowflake.snowpark_connect.utils.udf_utils as udf_utils
-from snowflake.snowpark import Column, Session
+from snowflake.snowpark import Session
 from snowflake.snowpark.types import DataType, _parse_datatype_json_value
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
 from snowflake.snowpark_connect.config import global_config
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
 from snowflake.snowpark_connect.expression.map_expression import (
     map_single_column_expression,
 )
+from snowflake.snowpark_connect.expression.map_unresolved_star import (
+    map_unresolved_star_as_single_column,
+)
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
+from snowflake.snowpark_connect.typed_column import TypedColumn
 from snowflake.snowpark_connect.utils.context import (
     get_is_aggregate_function,
     get_is_evaluating_join_condition,
@@ -38,6 +44,7 @@ class SnowparkUDF(NamedTuple):
     return_type: DataType
     input_types: list[DataType]
     original_return_type: DataType | None
+    cast_to_original_return_type: bool = False
 
 
 def require_creating_udf_in_sproc(
@@ -184,6 +191,7 @@ def parse_return_type(return_type_json_str) -> Optional[DataType]:
 
 
 def create(session, called_from, return_type_json_str, input_types_json_str, input_column_names_json_str, udf_name, replace, udf_packages, udf_imports, b64_str, original_return_type):
+    session._use_scoped_temp_objects = False
     import snowflake.snowpark.context as context
     context._use_structured_type_semantics = True
     context._is_snowpark_connect_compatible_mode = True
@@ -227,25 +235,15 @@ def _check_supported_udf(
         case "python_udf":
             pass
         case "java_udf":
-            from snowflake.snowpark_connect.utils.session import (
-                get_or_create_snowpark_session,
-            )
-
-            session = get_or_create_snowpark_session()
-            if udf_proto.java_udf.class_name not in session._cached_java_udfs:
-                raise AnalysisException(
-                    f"Can not load class {udf_proto.java_udf.class_name}"
-                )
-            else:
-                raise ValueError(
-                    "Function type java_udf not supported for common inline user-defined function"
-                )
+            pass
         case "scalar_scala_udf":
             pass
         case _ as function_type:
-            raise ValueError(
+            exception = ValueError(
                 f"Function type {function_type} not supported for common inline user-defined function"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def _aggregate_function_check(
@@ -253,9 +251,11 @@ def _aggregate_function_check(
 ):
     name, is_aggregate_function = get_is_aggregate_function()
     if not udf_proto.deterministic and name != "default" and is_aggregate_function:
-        raise AnalysisException(
+        exception = AnalysisException(
             f"[AGGREGATE_FUNCTION_WITH_NONDETERMINISTIC_EXPRESSION] Non-deterministic expression {name}({udf_proto.function_name}) should not appear in the arguments of an aggregate function."
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+        raise exception
 
 
 def _join_checks(snowpark_udf_arg_names: list[str]):
@@ -282,49 +282,51 @@ def _join_checks(snowpark_udf_arg_names: list[str]):
             and is_left_evaluable
             and is_right_evaluable
         ):
-            raise AnalysisException(
+            exception = AnalysisException(
                 f"Detected implicit cartesian product for {is_evaluating_join_condition[0]} join between logical plans. \n"
                 f"Join condition is missing or trivial. \n"
                 f"Either: use the CROSS JOIN syntax to allow cartesian products between those relations, or; "
                 f"enable implicit cartesian products by setting the configuration variable spark.sql.crossJoin.enabled=True."
             )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+            raise exception
         if (
             is_evaluating_join_condition[0] != "INNER"
             and is_evaluating_join_condition[1]
             and is_left_evaluable
             and is_right_evaluable
         ):
-            raise AnalysisException(
+            exception = AnalysisException(
                 f"[UNSUPPORTED_FEATURE.PYTHON_UDF_IN_ON_CLAUSE] The feature is not supported: "
                 f"Python UDF in the ON clause of a {is_evaluating_join_condition[0]} JOIN. "
                 f"In case of an INNNER JOIN consider rewriting to a CROSS JOIN with a WHERE clause."
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def infer_snowpark_arguments(
     udf_proto: expressions_proto.CommonInlineUserDefinedFunction,
     column_mapping: ColumnNameMap,
     typer: ExpressionTyper,
-) -> tuple[list[str], list[Column]]:
-    snowpark_udf_args: list[Column] = []
+) -> tuple[list[str], list[TypedColumn]]:
+    snowpark_udf_args: list[TypedColumn] = []
     snowpark_udf_arg_names: list[str] = []
     for arg_exp in udf_proto.arguments:
-        (
-            snowpark_udf_arg_name,
-            snowpark_udf_arg,
-        ) = map_single_column_expression(arg_exp, column_mapping, typer)
-        snowpark_udf_arg = snowpark_udf_arg.col
-        snowpark_udf_args.append(snowpark_udf_arg)
-        snowpark_udf_arg_names.append(snowpark_udf_arg_name)
+        # Handle unresolved_star expressions specially
+        if arg_exp.HasField("unresolved_star"):
+            # Use map_unresolved_star_as_struct to expand star into a single combined column
+            spark_name, typed_column = map_unresolved_star_as_single_column(
+                arg_exp, column_mapping, typer
+            )
+            snowpark_udf_args.append(typed_column)
+            snowpark_udf_arg_names.append(spark_name)
+        else:
+            (
+                snowpark_udf_arg_name,
+                snowpark_udf_arg,
+            ) = map_single_column_expression(arg_exp, column_mapping, typer)
+            snowpark_udf_args.append(snowpark_udf_arg)
+            snowpark_udf_arg_names.append(snowpark_udf_arg_name)
     _join_checks(snowpark_udf_arg_names)
     return snowpark_udf_arg_names, snowpark_udf_args
-
-
-def gen_input_types(
-    snowpark_udf_args: list[Column],
-    typer: ExpressionTyper,
-):
-    input_types = []
-    for udf_arg in snowpark_udf_args:
-        input_types.extend(typer.type(udf_arg))
-    return input_types

snowflake/snowpark_connect/utils/udf_utils.py

@@ -103,7 +103,7 @@ class ProcessCommonInlineUserDefinedFunction:
                 )
             case _:
                 raise ValueError(
-                    f"Function type {self._function_type} not supported for common inline user-defined function"
+                    f"[snowpark_connect::unsupported_operation] Function type {self._function_type} not supported for common inline user-defined function"
                 )
 
     @property
@@ -112,7 +112,7 @@ class ProcessCommonInlineUserDefinedFunction:
             return self._snowpark_udf_args
         else:
             raise ValueError(
-                "Column mapping is not provided, cannot get snowpark udf args"
+                "[snowpark_connect::internal_error] Column mapping is not provided, cannot get snowpark udf args"
             )
 
     @property
@@ -121,7 +121,7 @@ class ProcessCommonInlineUserDefinedFunction:
             return self._snowpark_udf_arg_names
         else:
             raise ValueError(
-                "Column mapping is not provided, cannot get snowpark udf arg names"
+                "[snowpark_connect::internal_error] Column mapping is not provided, cannot get snowpark udf arg names"
             )
 
     def _create_python_udf(self):
@@ -148,7 +148,12 @@ class ProcessCommonInlineUserDefinedFunction:
 
             # Change directory to the one containing the UDF imported files
             import_path = sys._xoptions["snowflake_import_directory"]
-            tmp_path = f"/tmp/sas-{os.getpid()}"
+            if os.name == "nt":
+                import tempfile
+
+                tmp_path = os.path.join(tempfile.gettempdir(), f"sas-{os.getpid()}")
+            else:
+                tmp_path = f"/tmp/sas-{os.getpid()}"
             os.makedirs(tmp_path, exist_ok=True)
             os.chdir(tmp_path)
             shutil.copytree(import_path, tmp_path, dirs_exist_ok=True)
@@ -176,14 +181,6 @@ class ProcessCommonInlineUserDefinedFunction:
                         tar_ref.extractall(archive[: -len(".archive")])
                 os.remove(archive)
 
-        def callable_func(*args, **kwargs):
-            import_staged_files()
-            return original_callable(*args, **kwargs)
-
-        callable_func.__signature__ = inspect.signature(original_callable)
-        if hasattr(original_callable, "__annotations__"):
-            callable_func.__annotations__ = original_callable.__annotations__
-
         if self._udf_packages:
             packages = [p.strip() for p in self._udf_packages.strip("[]").split(",")]
         else:
@@ -193,13 +190,109 @@ class ProcessCommonInlineUserDefinedFunction:
         else:
             imports = []
 
+        def callable_func(*args, **kwargs):
+            if imports:
+                import_staged_files()
+            return original_callable(*args, **kwargs)
+
+        callable_func.__signature__ = inspect.signature(original_callable)
+        if hasattr(original_callable, "__annotations__"):
+            callable_func.__annotations__ = original_callable.__annotations__
+
         update_none_input_types()
 
+        struct_positions = [
+            i
+            for i, t in enumerate(self._input_types or [])
+            if isinstance(t, StructType)
+        ]
+
+        if struct_positions:
+
+            class StructRowProxy:
+                """Row-like object supporting positional and named access for PySpark compatibility."""
+
+                def __init__(self, fields, values) -> None:
+                    self._fields = fields
+                    self._values = values
+                    self._field_to_index = {field: i for i, field in enumerate(fields)}
+
+                def __getitem__(self, key):
+                    if isinstance(key, int):
+                        return self._values[key]
+                    elif isinstance(key, str):
+                        if key in self._field_to_index:
+                            return self._values[self._field_to_index[key]]
+                        raise KeyError(f"Field '{key}' not found in struct")
+                    else:
+                        raise TypeError(f"Invalid key type: {type(key)}")
+
+                def __getattr__(self, name):
+                    if name.startswith("_"):
+                        raise AttributeError(f"Attribute '{name}' not found")
+                    if name in self._field_to_index:
+                        return self._values[self._field_to_index[name]]
+                    raise AttributeError(f"Attribute '{name}' not found")
+
+                def __len__(self):
+                    return len(self._values)
+
+                def __iter__(self):
+                    return iter(self._values)
+
+                def __repr__(self):
+                    field_values = [
+                        f"{field}={repr(value)}"
+                        for field, value in zip(self._fields, self._values)
+                    ]
+                    return f"Row({', '.join(field_values)})"
+
+                def asDict(self):
+                    """Convert to dict (like PySpark Row.asDict())."""
+                    return dict(zip(self._fields, self._values))
+
+            def convert_to_row(arg):
+                """Convert dict to StructRowProxy. Only called for struct positions."""
+                if isinstance(arg, dict) and arg:
+                    fields = list(arg.keys())
+                    values = [arg[k] for k in fields]
+                    return StructRowProxy(fields, values)
+                return arg
+
+            def convert_from_row(result):
+                """Convert StructRowProxy back to dict for serialization."""
+                if isinstance(result, StructRowProxy):
+                    return result.asDict()
+                return result
+
+        def struct_input_wrapper(*args, **kwargs):
+            if struct_positions:
+                processed_args = []
+                for i, arg in enumerate(args):
+                    if i in struct_positions:
+                        processed_args.append(convert_to_row(arg))
+                    else:
+                        processed_args.append(arg)
+
+                processed_kwargs = {k: convert_to_row(v) for k, v in kwargs.items()}
+                result = callable_func(*tuple(processed_args), **processed_kwargs)
+                # Convert any StructRowProxy in return value back to dict for serialization
+                return convert_from_row(result)
+            return callable_func(*args, **kwargs)
+
         needs_struct_conversion = isinstance(self._original_return_type, StructType)
 
+        # Use callable_func directly when there are no struct inputs to avoid closure issues.
+        # struct_input_wrapper captures convert_to_row in its closure, but convert_to_row is only
+        # defined when struct_positions is truthy. Cloudpickle serializes all closure variables,
+        # so using struct_input_wrapper without struct positions would fail during serialization.
+        updated_callable_func = (
+            struct_input_wrapper if struct_positions else callable_func
+        )
+
         if not needs_struct_conversion:
             return snowpark_fn.udf(
-                create_null_safe_wrapper(callable_func),
+                create_null_safe_wrapper(updated_callable_func),
                 return_type=self._return_type,
                 input_types=self._input_types,
                 name=self._udf_name,
@@ -225,7 +318,21 @@ class ProcessCommonInlineUserDefinedFunction:
         field_names = [field.name for field in self._original_return_type.fields]
 
         def struct_wrapper(*args):
+            if struct_positions:
+                processed_args = []
+                for i, arg in enumerate(args):
+                    if i in struct_positions:
+                        processed_args.append(convert_to_row(arg))
+                    else:
+                        processed_args.append(arg)
+                args = tuple(processed_args)
+
             result = callable_func(*args)
+
+            # Convert StructRowProxy back to dict for serialization
+            if struct_positions:
+                result = convert_from_row(result)
+
             if isinstance(result, (tuple, list)):
                 # Convert tuple/list to dict using struct field names
                 if len(result) == len(field_names):
@@ -283,6 +390,18 @@ class ProcessCommonInlineUserDefinedFunction:
             case "python_udf":
                 return self._create_python_udf()
             case "scalar_scala_udf":
+                from snowflake.snowpark_connect.utils.context import (
+                    get_is_aggregate_function,
+                )
+
+                name, is_aggregate_function = get_is_aggregate_function()
+                if is_aggregate_function and name.lower() == "reduce":
+                    # Handling of Scala Reduce function requires usage of Java UDAF
+                    from snowflake.snowpark_connect.utils.java_udaf_utils import (
+                        create_java_udaf_for_reduce_scala_function,
+                    )
+
+                    return create_java_udaf_for_reduce_scala_function(self)
                 from snowflake.snowpark_connect.utils.scala_udf_utils import (
                     create_scala_udf,
                 )
@@ -290,5 +409,5 @@ class ProcessCommonInlineUserDefinedFunction:
                 return create_scala_udf(self)
             case _:
                 raise ValueError(
-                    f"Function type {self._function_type} not supported for common inline user-defined function"
+                    f"[snowpark_connect::unsupported_operation] Function type {self._function_type} not supported for common inline user-defined function"
                 )

snowflake/snowpark_connect/utils/udtf_helper.py

@@ -16,6 +16,8 @@ import snowflake.snowpark_connect.tcm as tcm
 from snowflake import snowpark
 from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted
 from snowflake.snowpark.types import DataType, StructType, _parse_datatype_json_value
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
 from snowflake.snowpark_connect.type_mapping import proto_to_snowpark_type
 from snowflake.snowpark_connect.utils import pandas_udtf_utils, udtf_utils
 from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
@@ -37,7 +39,9 @@ def udtf_check(
     udtf_proto: relation_proto.CommonInlineUserDefinedTableFunction,
 ) -> None:
     if udtf_proto.WhichOneof("function") != "python_udtf":
-        raise ValueError(f"Not python udtf {udtf_proto.function}")
+        exception = ValueError(f"Not python udtf {udtf_proto.function}")
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
 
 def require_creating_udtf_in_sproc(
@@ -149,6 +153,7 @@ def parse_types(types_json_str) -> Optional[list[DataType]]:
     return json.loads(types_json_str)
 
 def create(session, b64_str, expected_types_json_str, output_schema_json_str, packages, imports, is_arrow_enabled, is_spark_compatible_udtf_mode_enabled, called_from):
+    session._use_scoped_temp_objects = False
     import snowflake.snowpark.context as context
     context._use_structured_type_semantics = True
     context._is_snowpark_connect_compatible_mode = True
@@ -253,6 +258,7 @@ from snowflake.snowpark.types import _parse_datatype_json_value
 {inline_udtf_utils_py_code}
 
 def create(session, b64_str, spark_column_names_json_str, input_schema_json_str, return_schema_json_str):
+    session._use_scoped_temp_objects = False
     import snowflake.snowpark.context as context
     context._use_structured_type_semantics = True
     context._is_snowpark_connect_compatible_mode = True
@@ -326,6 +332,7 @@ from snowflake.snowpark.types import _parse_datatype_json_value
 from pyspark.serializers import CloudPickleSerializer
 
 def create(session, func_info_json):
+    session._use_scoped_temp_objects = False
     import snowflake.snowpark.context as context
     context._use_structured_type_semantics = True
     context._is_snowpark_connect_compatible_mode = True

snowflake/snowpark_connect/utils/udtf_utils.py

@@ -32,14 +32,16 @@ def create_udtf(
     udtf = udtf_proto.python_udtf
     callable_func = CloudPickleSerializer().loads(udtf.command)
 
-    func_signature = inspect.signature(callable_func.eval)
+    original_func = callable_func.eval
+    func_signature = inspect.signature(original_func)
     # Set all input types to VariantType regardless of type hints so that we can pass all arguments as VariantType.
     # Otherwise, we will run into issues with type mismatches. This only applies for UDTF registration.
     # We subtract one here since UDTF functions are class methods and always have "self" as the first parameter.
     input_types = [VariantType()] * (len(func_signature.parameters) - 1)
 
-    # Wrapp callable to allow reading imported files
-    callable_func = artifacts_reader_wrapper(callable_func)
+    if imports:
+        # Wrapp callable to allow reading imported files
+        callable_func = artifacts_reader_wrapper(callable_func)
 
     if is_arrow_enabled:
         callable_func = spark_compatible_udtf_wrapper_with_arrow(
@@ -48,7 +50,7 @@ def create_udtf(
     elif is_spark_compatible_udtf_mode_enabled:
         callable_func = spark_compatible_udtf_wrapper(callable_func, expected_types)
     else:
-        callable_func.process = callable_func.eval
+        callable_func.process = original_func
         if hasattr(callable_func, "terminate"):
             callable_func.end_partition = callable_func.terminate
 
@@ -107,7 +109,9 @@ def create_udtf(
                 imports=imports,
             )
         case _:
-            raise NotImplementedError(f"{called_from}")
+            raise NotImplementedError(
+                f"[snowpark_connect::unsupported_operation] {called_from}"
+            )
 
 
 def artifacts_reader_wrapper(user_udtf_cls: type) -> type:
@@ -127,7 +131,12 @@ def artifacts_reader_wrapper(user_udtf_cls: type) -> type:
 
                 # Change directory to the one containing the UDF imported files
                 import_path = sys._xoptions["snowflake_import_directory"]
-                tmp_path = f"/tmp/sas-{os.getpid()}"
+                if os.name == "nt":
+                    import tempfile
+
+                    tmp_path = os.path.join(tempfile.gettempdir(), f"sas-{os.getpid()}")
+                else:
+                    tmp_path = f"/tmp/sas-{os.getpid()}"
                 os.makedirs(tmp_path, exist_ok=True)
                 os.chdir(tmp_path)
                 shutil.copytree(import_path, tmp_path, dirs_exist_ok=True)
@@ -195,17 +204,19 @@ def _create_convert_table_argument_to_row():
                 # Named access: row["col1"], row["col2"]
                 if key in self._field_to_index:
                     return self._values[self._field_to_index[key]]
-                raise KeyError(key)
+                raise KeyError(f"[snowpark_connect::invalid_operation] {key}")
             else:
-                raise TypeError(f"Invalid key type: {type(key)}")
+                raise TypeError(
+                    f"[snowpark_connect::type_mismatch] Invalid key type: {type(key)}"
+                )
 
         def __getattr__(self, name):
             # Attribute access: row.col1, row.col2
             if name.startswith("_"):
-                raise AttributeError(name)
+                raise AttributeError(f"[snowpark_connect::invalid_operation] {name}")
             if name in self._field_to_index:
                 return self._values[self._field_to_index[name]]
-            raise AttributeError(name)
+            raise AttributeError(f"[snowpark_connect::invalid_operation] {name}")
 
         def __len__(self):
             return len(self._values)
@@ -279,7 +290,9 @@ def spark_compatible_udtf_wrapper(
             return val
         if isinstance(val, datetime.datetime):
             return val.date()
-        raise AttributeError(f"Invalid date value {val}")
+        raise AttributeError(
+            f"[snowpark_connect::invalid_input] Invalid date value {val}"
+        )
 
     def _coerce_to_binary(val: object, target_type_name: str = "byte") -> bytes | None:
         if target_type_name == "binary":
@@ -343,7 +356,9 @@ def spark_compatible_udtf_wrapper(
     def _coerce_to_timestamp(val: object) -> datetime.datetime | None:
         if isinstance(val, datetime.datetime):
             return val
-        raise AttributeError(f"Invalid time stamp value {val}")
+        raise AttributeError(
+            f"[snowpark_connect::invalid_input] Invalid time stamp value {val}"
+        )
 
     SCALAR_COERCERS = {
         "bool": _coerce_to_bool,
@@ -447,7 +462,7 @@ def spark_compatible_udtf_wrapper(
 
                 if not isinstance(raw_row_tuple, (tuple, list)):
                     raise TypeError(
-                        f"[UDTF_INVALID_OUTPUT_ROW_TYPE] return value should be an iterable object containing tuples, but got {type(raw_row_tuple)}"
+                        f"[snowpark_connect::type_mismatch] [UDTF_INVALID_OUTPUT_ROW_TYPE] return value should be an iterable object containing tuples, but got {type(raw_row_tuple)}"
                     )
 
                 if len(raw_row_tuple) != len(expected_types):
@@ -467,7 +482,7 @@ def spark_compatible_udtf_wrapper(
                         and val is not None
                     ):
                         raise RuntimeError(
-                            f"[UNEXPECTED_TUPLE_WITH_STRUCT] Expected a struct for column at position {i}, but got a primitive value of type {type(val)}"
+                            f"[snowpark_connect::type_mismatch] [UNEXPECTED_TUPLE_WITH_STRUCT] Expected a struct for column at position {i}, but got a primitive value of type {type(val)}"
                         )
 
                 coerced_row_list = [None] * len(expected_types)
@@ -533,7 +548,7 @@ def spark_compatible_udtf_wrapper_with_arrow(
                     return pa.map_(key_type, value_type)
                 case _, _:
                     raise TypeError(
-                        f"[UDTF_ARROW_TYPE_CAST_ERROR] Unsupported Python scalar type for Arrow conversion: {target_py_type}"
+                        f"[snowpark_connect::unsupported_type] [UDTF_ARROW_TYPE_CAST_ERROR] Unsupported Python scalar type for Arrow conversion: {target_py_type}"
                     )
         elif kind == "array":
             element_type_info = type_marker
@@ -543,7 +558,7 @@ def spark_compatible_udtf_wrapper_with_arrow(
             struct_fields_info = type_marker
             if not isinstance(struct_fields_info, dict):
                 raise TypeError(
-                    f"[UDTF_ARROW_TYPE_CAST_ERROR] Invalid struct definition for Arrow: expected dict, got {type(struct_fields_info)}"
+                    f"[snowpark_connect::type_mismatch] [UDTF_ARROW_TYPE_CAST_ERROR] Invalid struct definition for Arrow: expected dict, got {type(struct_fields_info)}"
                 )
             fields = []
             for field_name, field_type_info in struct_fields_info.items():
@@ -552,7 +567,7 @@ def spark_compatible_udtf_wrapper_with_arrow(
             return pa.struct(fields)
         else:
             raise TypeError(
-                f"[UDTF_ARROW_TYPE_CAST_ERROR] Unsupported data kind for Arrow conversion: {kind}"
+                f"[snowpark_connect::unsupported_type] [UDTF_ARROW_TYPE_CAST_ERROR] Unsupported data kind for Arrow conversion: {kind}"
             )
 
     def _convert_to_arrow_value(
@@ -576,7 +591,7 @@ def spark_compatible_udtf_wrapper_with_arrow(
                 ]
             if not isinstance(obj, (list, tuple)):
                 raise TypeError(
-                    f"[UDTF_ARROW_TYPE_CAST_ERROR] Expected list or tuple for Arrow array type, got {type(obj).__name__}"
+                    f"[snowpark_connect::type_mismatch] [UDTF_ARROW_TYPE_CAST_ERROR] Expected list or tuple for Arrow array type, got {type(obj).__name__}"
                 )
             element_type = arrow_type.value_type
             return [_convert_to_arrow_value(e, element_type, "array") for e in obj]
@@ -584,7 +599,7 @@ def spark_compatible_udtf_wrapper_with_arrow(
         if pa.types.is_map(arrow_type):
             if not isinstance(obj, dict):
                 raise TypeError(
-                    f"[UDTF_ARROW_TYPE_CAST_ERROR] Expected dict for Arrow map type, got {type(obj).__name__}"
+                    f"[snowpark_connect::type_mismatch] [UDTF_ARROW_TYPE_CAST_ERROR] Expected dict for Arrow map type, got {type(obj).__name__}"
                 )
             key_type = arrow_type.key_type
             value_type = arrow_type.item_type
@@ -610,7 +625,7 @@ def spark_compatible_udtf_wrapper_with_arrow(
             else:
                 # If the UDTF yields a list/tuple (or anything not a dict) for a struct column, it's an error.
                 raise TypeError(
-                    f"[UDTF_ARROW_TYPE_CAST_ERROR] Expected a dictionary for Arrow struct type column, but got {type(obj).__name__}"
+                    f"[snowpark_connect::type_mismatch] [UDTF_ARROW_TYPE_CAST_ERROR] Expected a dictionary for Arrow struct type column, but got {type(obj).__name__}"
                 )
 
         # Check if a scalar type is expected and if obj is a collection; if so, error out.
@@ -622,7 +637,7 @@ def spark_compatible_udtf_wrapper_with_arrow(
         ):
             if isinstance(obj, (list, tuple, dict)):
                 raise TypeError(
-                    f"[UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert Python collection type {type(obj).__name__} to scalar Arrow type {arrow_type}"
+                    f"[snowpark_connect::type_mismatch] [UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert Python collection type {type(obj).__name__} to scalar Arrow type {arrow_type}"
                 )
 
         if pa.types.is_boolean(arrow_type):
@@ -638,7 +653,7 @@ def spark_compatible_udtf_wrapper_with_arrow(
                 elif obj == 1:
                     return True
                 raise TypeError(
-                    f"[UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert {obj} to Arrow boolean"
+                    f"[snowpark_connect::type_mismatch] [UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert {obj} to Arrow boolean"
                 )
             if isinstance(obj, str):
                 v_str = obj.strip().lower()
@@ -647,7 +662,7 @@ def spark_compatible_udtf_wrapper_with_arrow(
                 if v_str == "false":
                     return False
             raise TypeError(
-                f"[UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert {type(obj).__name__} to Arrow boolean"
+                f"[snowpark_connect::type_mismatch] [UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert {type(obj).__name__} to Arrow boolean"
             )
 
         if pa.types.is_integer(arrow_type):
@@ -663,7 +678,7 @@ def spark_compatible_udtf_wrapper_with_arrow(
                 except ValueError:
                     pass
             raise TypeError(
-                f"[UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert {type(obj).__name__} to Arrow integer"
+                f"[snowpark_connect::type_mismatch] [UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert {type(obj).__name__} to Arrow integer"
             )
 
         if pa.types.is_floating(arrow_type):
@@ -675,7 +690,7 @@ def spark_compatible_udtf_wrapper_with_arrow(
                 except ValueError:
                     pass
             raise TypeError(
-                f"[UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert {type(obj).__name__} to Arrow float"
+                f"[snowpark_connect::type_mismatch] [UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert {type(obj).__name__} to Arrow float"
             )
 
         if pa.types.is_string(arrow_type):
@@ -687,7 +702,7 @@ def spark_compatible_udtf_wrapper_with_arrow(
             if isinstance(obj, str):
                 return obj
             raise TypeError(
-                f"[UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert {type(obj).__name__} to Arrow string"
+                f"[snowpark_connect::type_mismatch] [UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert {type(obj).__name__} to Arrow string"
             )
 
         if pa.types.is_binary(arrow_type) or pa.types.is_fixed_size_binary(arrow_type):
@@ -698,21 +713,21 @@ def spark_compatible_udtf_wrapper_with_arrow(
             if isinstance(obj, int):
                 return bytearray([obj])
             raise TypeError(
-                f"[UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert {type(obj).__name__} to Arrow binary"
+                f"[snowpark_connect::type_mismatch] [UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert {type(obj).__name__} to Arrow binary"
             )
 
         if pa.types.is_date(arrow_type):
             if isinstance(obj, datetime.date):
                 return obj
             raise TypeError(
-                f"[UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert {type(obj).__name__} to Arrow date. Expected datetime.date."
+                f"[snowpark_connect::type_mismatch] [UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert {type(obj).__name__} to Arrow date. Expected datetime.date."
             )
 
         if pa.types.is_timestamp(arrow_type):
             if isinstance(obj, datetime.datetime):
                 return obj
             raise TypeError(
-                f"[UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert {type(obj).__name__} to Arrow timestamp. Expected datetime.datetime."
+                f"[snowpark_connect::type_mismatch] [UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert {type(obj).__name__} to Arrow timestamp. Expected datetime.datetime."
             )
 
         if pa.types.is_decimal(arrow_type):
@@ -727,11 +742,11 @@ def spark_compatible_udtf_wrapper_with_arrow(
                     pass
 
             raise TypeError(
-                f"[UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert {type(obj).__name__} to Arrow decimal. Expected decimal.Decimal or compatible int/str."
+                f"[snowpark_connect::type_mismatch] [UDTF_ARROW_TYPE_CAST_ERROR] Cannot convert {type(obj).__name__} to Arrow decimal. Expected decimal.Decimal or compatible int/str."
             )
 
         raise TypeError(
-            f"[UDTF_ARROW_TYPE_CAST_ERROR] Unsupported type conversion for {type(obj).__name__} to Arrow type {arrow_type}"
+            f"[snowpark_connect::unsupported_operation] [UDTF_ARROW_TYPE_CAST_ERROR] Unsupported type conversion for {type(obj).__name__} to Arrow type {arrow_type}"
         )
 
     class WrappedUDTF: