snowpark-connect 0.30.1__py3-none-any.whl → 0.31.0__py3-none-any.whl

snowflake/snowpark_connect/column_name_handler.py

@@ -20,7 +20,12 @@ 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.utils.context import get_current_operation_scope
+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.utils.context import (
+    get_current_operation_scope,
+    get_is_processing_order_by,
+)
 from snowflake.snowpark_connect.utils.identifiers import (
     split_fully_qualified_spark_name,
 )
@@ -142,7 +147,7 @@ class ColumnNameMap:
         parent_column_name_map: parent ColumnNameMap
         """
         self.columns: list[ColumnNames] = []
-        self.spark_to_col = defaultdict(list)
+        self.spark_to_col: defaultdict[str, list[ColumnNames]] = defaultdict(list)
         self.uppercase_spark_to_col = defaultdict(list)
         self.snowpark_to_col = defaultdict(list)
         self.is_case_sensitive = is_case_sensitive
@@ -353,18 +358,100 @@ class ColumnNameMap:
 
         snowpark_names_len = len(snowpark_names)
         if snowpark_names_len > 1:
-            raise AnalysisException(
-                f"Ambiguous spark column name {spark_column_name}, potential snowpark column names {snowpark_names}"
-            )
+            # Check if this is a case where we have identical expressions that can be safely resolved to the first one
+            # This commonly happens with GROUP BY expressions that also appear in SELECT clauses
+            if (
+                get_is_processing_order_by()
+                and self._can_resolve_ambiguous_identical_expressions(
+                    resolved_name, snowpark_names
+                )
+            ):
+                # All the ambiguous columns represent the same expression, so we can safely use the first one
+                return snowpark_names[0]
+            else:
+                exception = AnalysisException(
+                    f"Ambiguous spark column name {spark_column_name}, potential snowpark column names {snowpark_names}"
+                )
+                attach_custom_error_code(exception, ErrorCodes.AMBIGUOUS_COLUMN_NAME)
+                raise exception
         elif snowpark_names_len == 0:
             if allow_non_exists:
                 return None
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Spark column name {spark_column_name} does not exist"
                 )
+                attach_custom_error_code(exception, ErrorCodes.COLUMN_NOT_FOUND)
+                raise exception
         return snowpark_names[0]
 
+    def _can_resolve_ambiguous_identical_expressions(
+        self, spark_column_name: str, snowpark_names: list[str]
+    ) -> bool:
+        """
+        Determine if ambiguous columns represent identical expressions that can be safely resolved to the first one.
+
+        This handles the common case where the same expression (like a UDF call) appears multiple times
+        in a SELECT clause within a GROUP BY query. Since they're the same expression operating on the
+        same grouped data, they will have identical values, so we can safely resolve to any of them.
+
+        Args:
+            spark_column_name: The Spark column name that has multiple mappings, make sure resolve this reforehand
+            snowpark_names: List of Snowpark column names that map to this Spark column name
+
+        Returns:
+            True if we can safely resolve to the first snowpark column, False otherwise
+        """
+        if spark_column_name not in self.spark_to_col:
+            return False
+
+        columns = self.spark_to_col[spark_column_name]
+
+        # If we don't have multiple columns, there's no ambiguity to resolve
+        if len(columns) <= 1:
+            return False
+
+        # Check if all the snowpark names correspond to columns that have identical underlying expressions
+        # We'll compare the actual column objects to see if they represent the same computation
+        first_column = columns[0]
+
+        for column in columns[1:]:
+            # Check snowpark_type attribute
+            # If one has the attribute but the other doesn't, they're different
+            if hasattr(first_column, "snowpark_type") != hasattr(
+                column, "snowpark_type"
+            ):
+                return False
+            # If both have the attribute and values differ, they're different expressions
+            if (
+                hasattr(first_column, "snowpark_type")
+                and hasattr(column, "snowpark_type")
+                and first_column.snowpark_type != column.snowpark_type
+            ):
+                return False
+
+            # Check qualifiers attribute
+            # If one has the attribute but the other doesn't, they're different
+            if hasattr(first_column, "qualifiers") != hasattr(column, "qualifiers"):
+                return False
+            # If both have the attribute and values differ, they might be from different contexts
+            if (
+                hasattr(first_column, "qualifiers")
+                and hasattr(column, "qualifiers")
+                and first_column.qualifiers != column.qualifiers
+            ):
+                return False
+
+        # Additional safety check: ensure all snowpark names are actually in our mapping
+        for snowpark_name in snowpark_names:
+            if snowpark_name not in self.snowpark_to_col:
+                return False
+
+        # If we reach here, the columns appear to be identical expressions from the same context
+        # This commonly happens in GROUP BY scenarios where the same expression appears in both
+        # the grouping clause and the select clause
+        return True
+
     def get_spark_column_names_from_snowpark_column_names(
         self,
         snowpark_column_names: list[str],
@@ -390,16 +477,20 @@ class ColumnNameMap:
         )
         spark_names_len = len(spark_names)
         if spark_names_len > 1:
-            raise AnalysisException(
+            exception = AnalysisException(
                 f"Ambiguous snowpark column name {snowpark_column_name}, potential spark column names {spark_names}"
             )
+            attach_custom_error_code(exception, ErrorCodes.AMBIGUOUS_COLUMN_NAME)
+            raise exception
         elif spark_names_len == 0:
             if allow_non_exists:
                 return None
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Snowpark column name {snowpark_column_name} does not exist"
                 )
+                attach_custom_error_code(exception, ErrorCodes.COLUMN_NOT_FOUND)
+                raise exception
         return spark_names[0]
 
     def get_spark_column_name(self, idx: int) -> str:
@@ -604,14 +695,18 @@ class JoinColumnNameMap(ColumnNameMap):
             if allow_non_exists:
                 return None
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Spark column name {spark_column_name} does not exist in either left or right DataFrame"
                 )
+                attach_custom_error_code(exception, ErrorCodes.COLUMN_NOT_FOUND)
+                raise exception
 
         if (snowpark_column_name_in_right is not None) and (
             snowpark_column_name_in_left is not None
         ):
-            raise AnalysisException(f"Ambiguous column name {spark_column_name}")
+            exception = AnalysisException(f"Ambiguous column name {spark_column_name}")
+            attach_custom_error_code(exception, ErrorCodes.AMBIGUOUS_COLUMN_NAME)
+            raise exception
 
         snowpark_name = (
             snowpark_column_name_in_right
@@ -637,57 +732,85 @@ class JoinColumnNameMap(ColumnNameMap):
     def get_snowpark_column_names_from_spark_column_names(
         self, spark_column_names: list[str], return_first: bool = False
     ) -> list[str]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_spark_column_names_from_snowpark_column_names(
         self,
         snowpark_column_names: list[str],
     ) -> list[str]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_spark_column_name_from_snowpark_column_name(
         self, snowpark_column_name: str
     ) -> str:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_spark_columns(self) -> list[str]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_snowpark_columns(self) -> list[str]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_snowpark_columns_after_drop(self, cols_to_drop: list[str]) -> list[str]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_renamed_nested_column_name(self, name) -> str | None:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def has_spark_column(self, spark_column_name: str) -> bool:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def snowpark_to_spark_map(self) -> dict[str, str]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def spark_to_snowpark_for_pattern(self, pattern: str) -> list[tuple[str, str]]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def with_columns(
         self, new_spark_columns: list[str], new_snowpark_columns: list[str]
     ) -> tuple[list[str], list[str], list[list[str]]]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_qualifiers(self) -> list[list[str]]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_qualifiers_for_columns_after_drop(
         self, cols_to_drop: list[str]
     ) -> list[list[str]]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_spark_and_snowpark_columns_with_qualifier_for_qualifier(
         self, qualifiers_input: list[str]
     ) -> tuple[list[str], list[str], list[list[str]]]:
-        raise NotImplementedError("Method not implemented!")
+        exception = NotImplementedError("Method not implemented!")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     def get_qualifier_for_spark_column(self, spark_column_name: str) -> list[str]:
 
@@ -699,6 +822,8 @@ class JoinColumnNameMap(ColumnNameMap):
         )
 
         if (len(qualifier_left) > 0) and (len(qualifier_right) > 0):
-            raise AnalysisException(f"Ambiguous column name {spark_column_name}")
+            exception = AnalysisException(f"Ambiguous column name {spark_column_name}")
+            attach_custom_error_code(exception, ErrorCodes.AMBIGUOUS_COLUMN_NAME)
+            raise exception
 
         return qualifier_right if len(qualifier_left) == 0 else qualifier_left

snowflake/snowpark_connect/config.py

@@ -21,8 +21,13 @@ from snowflake.snowpark._internal.analyzer.analyzer_utils import (
 )
 from snowflake.snowpark.exceptions import SnowparkSQLException
 from snowflake.snowpark.types import TimestampTimeZone, TimestampType
+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.utils.concurrent import SynchronizedDict
-from snowflake.snowpark_connect.utils.context import get_session_id
+from snowflake.snowpark_connect.utils.context import (
+    get_jpype_jclass_lock,
+    get_session_id,
+)
 from snowflake.snowpark_connect.utils.external_udxf_cache import (
     clear_external_udxf_cache,
 )
@@ -291,7 +296,6 @@ class SessionConfig:
     """This class contains the session configuration for the Spark Server."""
 
     default_session_config = {
-        "snowpark.connect.sql.identifiers.auto-uppercase": "all_except_columns",
         "snowpark.connect.sql.passthrough": "false",
         "snowpark.connect.cte.optimization_enabled": "false",
         "snowpark.connect.udtf.compatibility_mode": "false",
@@ -359,9 +363,11 @@ def route_config_proto(
                 if not pair.HasField("value"):
                     from pyspark.errors import IllegalArgumentException
 
-                    raise IllegalArgumentException(
+                    exception = IllegalArgumentException(
                         f"Cannot set config '{pair.key}' to None"
                     )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_CONFIG_VALUE)
+                    raise exception
 
                 set_config_param(
                     config.session_id, pair.key, pair.value, snowpark_session
@@ -444,7 +450,11 @@ def route_config_proto(
                 pair.value = str(global_config.is_modifiable(key)).lower()
             return res
         case _:
-            raise SnowparkConnectNotImplementedError(f"Unexpected request {config}")
+            exception = SnowparkConnectNotImplementedError(
+                f"Unexpected request {config}"
+            )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def set_config_param(
@@ -484,19 +494,27 @@ def _verify_static_config_not_modified(key: str) -> None:
     # https://github.com/apache/spark/blob/v3.5.3/sql/core/src/main/scala/org/apache/spark/sql/RuntimeConfig.scala#L161
     # Spark does not allow to modify static configurations at runtime.
     if global_config.is_static_config(key) and global_config.is_set(key):
-        raise ValueError(f"Cannot modify the value of a static config: {key}")
+        exception = ValueError(f"Cannot modify the value of a static config: {key}")
+        attach_custom_error_code(exception, ErrorCodes.CONFIG_CHANGE_NOT_ALLOWED)
+        raise exception
 
 
 def _verify_is_valid_config_value(key: str, value: Any) -> None:
     if key in CONFIG_ALLOWED_VALUES and value not in CONFIG_ALLOWED_VALUES[key]:
-        raise ValueError(
+        exception = ValueError(
             f"Invalid value '{value}' for key '{key}'. Allowed values: {', '.join(CONFIG_ALLOWED_VALUES[key])}."
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_CONFIG_VALUE)
+        raise exception
 
 
 def _verify_is_not_readonly_config(key):
     if key in global_config.readonly_config_list:
-        raise ValueError(f"Config with key {key} is read-only and cannot be modified.")
+        exception = ValueError(
+            f"Config with key {key} is read-only and cannot be modified."
+        )
+        attach_custom_error_code(exception, ErrorCodes.CONFIG_CHANGE_NOT_ALLOWED)
+        raise exception
 
 
 def set_jvm_timezone(timezone_id: str):
@@ -513,10 +531,13 @@ def set_jvm_timezone(timezone_id: str):
         RuntimeError: If JVM is not started
     """
     if not jpype.isJVMStarted():
-        raise RuntimeError("JVM must be started before setting timezone")
+        exception = RuntimeError("JVM must be started before setting timezone")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     try:
-        TimeZone = jpype.JClass("java.util.TimeZone")
+        with get_jpype_jclass_lock():
+            TimeZone = jpype.JClass("java.util.TimeZone")
         new_timezone = TimeZone.getTimeZone(timezone_id)
         TimeZone.setDefault(new_timezone)
 
@@ -528,7 +549,9 @@ def set_jvm_timezone(timezone_id: str):
 def reset_jvm_timezone_to_system_default():
     """Reset JVM timezone to the system's default timezone"""
     if not jpype.isJVMStarted():
-        raise RuntimeError("JVM must be started first")
+        exception = RuntimeError("JVM must be started first")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     try:
         TimeZone = jpype.JClass("java.util.TimeZone")
@@ -537,9 +560,13 @@ def reset_jvm_timezone_to_system_default():
             f"Reset JVM timezone to system default: {TimeZone.getDefault().getID()}"
         )
     except jpype.JException as e:
-        raise RuntimeError(f"Java exception while resetting timezone: {e}")
+        exception = RuntimeError(f"Java exception while resetting timezone: {e}")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
     except Exception as e:
-        raise RuntimeError(f"Unexpected error resetting JVM timezone: {e}")
+        exception = RuntimeError(f"Unexpected error resetting JVM timezone: {e}")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
 
 def set_snowflake_parameters(
@@ -629,16 +656,24 @@ def get_describe_cache_ttl_seconds() -> int:
 
 def auto_uppercase_column_identifiers() -> bool:
     session_config = sessions_config[get_session_id()]
-    return session_config[
+    auto_upper_case_config = session_config[
         "snowpark.connect.sql.identifiers.auto-uppercase"
-    ].lower() in ("all", "only_columns")
+    ]
+    if auto_upper_case_config:
+        return auto_upper_case_config.lower() in ("all", "only_columns")
+
+    return not global_config.spark_sql_caseSensitive
 
 
 def auto_uppercase_non_column_identifiers() -> bool:
     session_config = sessions_config[get_session_id()]
-    return session_config[
+    auto_upper_case_config = session_config[
         "snowpark.connect.sql.identifiers.auto-uppercase"
-    ].lower() in ("all", "all_except_columns")
+    ]
+    if auto_upper_case_config:
+        return auto_upper_case_config.lower() in ("all", "all_except_columns")
+
+    return not global_config.spark_sql_caseSensitive
 
 
 def parse_imports(session: snowpark.Session, imports: str | None) -> None:

snowflake/snowpark_connect/date_time_format_mapping.py

@@ -24,6 +24,8 @@
 from pyspark.errors.exceptions.base import DateTimeException
 
 from snowflake.snowpark.types import DataType, StringType
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
 
 
 # TODO: There are more patterns where spark may throw an error.
@@ -229,9 +231,11 @@ def convert_spark_format_to_snowflake(
     timestamp_input_type: DataType | None = None,
 ):
     if spark_format in {"Y", "w", "W"}:
-        raise DateTimeException(
+        exception = DateTimeException(
             f"Fail to recognize '{spark_format}' pattern in the DateTimeFormatter."
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT)
+        raise exception
     snowflake_format = ""
     i = 0
     n = len(spark_format)
@@ -299,28 +303,46 @@ def convert_spark_format_to_snowflake(
                     # Spark's 'a' would be at most 1 times
                     is_valid_a_pattern = spark_format[i : i + 2] != char * 2
                     if not is_valid_a_pattern:
-                        raise DateTimeException(
+                        exception = DateTimeException(
                             f"Fail to recognize '{spark_format}' pattern in the DateTimeFormatter"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "h" | "K" | "k" | "H" | "m" | "s" | "d":
                     # Spark's characters would be at most 2 times
                     is_valid_2_patterns = spark_format[i : i + 3] != char * 3
                     if not is_valid_2_patterns:
-                        raise DateTimeException(
+                        exception = DateTimeException(
                             f"Fail to recognize '{spark_format}' pattern in the DateTimeFormatter"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "D":
                     # Spark's 'D'' would be at most 3 times
                     is_valid_D_patterns = spark_format[i : i + 4] != char * 4
                     if not is_valid_D_patterns:
-                        raise DateTimeException(
+                        exception = DateTimeException(
                             f"Fail to recognize '{spark_format}' pattern in the DateTimeFormatter"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "V":
                     # Spark's 'V' for Zone ID requires 'VV'. A single 'V' is invalid.
                     is_valid_vv_pattern = spark_format[i : i + 2] == "VV"
                     if not is_valid_vv_pattern:
-                        raise DateTimeException("Pattern letter count must be 2: V")
+                        exception = DateTimeException(
+                            "Pattern letter count must be 2: V"
+                        )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "O":
                     # Spark's 'O' would be either 1 or 4.
                     is_valid_o_or_oooo_pattern = spark_format[i : i + 2] != "OO" or (
@@ -328,28 +350,52 @@ def convert_spark_format_to_snowflake(
                         and spark_format[i : i + 5] != "OOOOO"
                     )
                     if not is_valid_o_or_oooo_pattern:
-                        raise DateTimeException(
+                        exception = DateTimeException(
                             "Pattern letter count must be 1 or 4: O"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "q" | "Q" | "z" | "E":
                     # Spark's characters would be at most 4 times
                     is_valid_4_patterns = spark_format[i : i + 5] != char * 5
                     if not is_valid_4_patterns:
-                        raise DateTimeException(f"Too many pattern letters: {char}")
+                        exception = DateTimeException(
+                            f"Too many pattern letters: {char}"
+                        )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "x" | "X" | "Z":
                     # Spark's 'x' or 'X' or 'z' or 'Z' would be at most 5 times
                     is_valid_xz_pattern = spark_format[i : i + 6] != char * 6
                     if not is_valid_xz_pattern:
-                        raise DateTimeException(f"Too many pattern letters: {char}")
+                        exception = DateTimeException(
+                            f"Too many pattern letters: {char}"
+                        )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "y":
                     # Spark's 'y' would be at most 6 times
                     is_valid_y_pattern = spark_format[i : i + 7] != char * 7
                     if not is_valid_y_pattern:
-                        raise DateTimeException(
+                        exception = DateTimeException(
                             f"Fail to recognize '{spark_format}' pattern in the DateTimeFormatter"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "C" | "I":
-                    raise DateTimeException(f"Unknown pattern letter: {char}")
+                    exception = DateTimeException(f"Unknown pattern letter: {char}")
+                    attach_custom_error_code(
+                        exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                    )
+                    raise exception
 
             if (
                 spark_format[i : i + 5] in {"M" * 5, "L" * 5}
@@ -362,9 +408,13 @@ def convert_spark_format_to_snowflake(
                 or spark_format[i : i + 3] in {"kkk", "KKK"}
                 or spark_format[i : i + 10] == "SSSSSSSSSS"
             ):
-                raise DateTimeException(
+                exception = DateTimeException(
                     f"Fail to recognize '{spark_format}' pattern in the DateTimeFormatter."
                 )
+                attach_custom_error_code(
+                    exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                )
+                raise exception
 
             matched_pattern = False
 
@@ -375,7 +425,11 @@ def convert_spark_format_to_snowflake(
                         spark_key
                     ]
                     if isinstance(snowflake_equivalent, _UnsupportedSparkFormatPattern):
-                        raise DateTimeException(snowflake_equivalent.message)
+                        exception = DateTimeException(snowflake_equivalent.message)
+                        attach_custom_error_code(
+                            exception, ErrorCodes.UNSUPPORTED_OPERATION
+                        )
+                        raise exception
                     if snowflake_equivalent is not None:
                         snowflake_format += snowflake_equivalent
                     i += len(spark_key)
@@ -389,7 +443,11 @@ def convert_spark_format_to_snowflake(
                     isinstance(timestamp_input_type, StringType)
                     and char not in snowflake_time_format_separator
                 ):
-                    raise DateTimeException(f"Illegal pattern character: {char}")
+                    exception = DateTimeException(f"Illegal pattern character: {char}")
+                    attach_custom_error_code(
+                        exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                    )
+                    raise exception
 
                 snowflake_format += f'"{char}"'
                 i += 1

snowflake/snowpark_connect/error/error_codes.py

@@ -0,0 +1,50 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+"""
+Error code constants for Snowpark Connect.
+
+This module defines custom error codes that can be attached to exceptions
+and included in gRPC error responses.
+"""
+
+
+class ErrorCodes:
+    """Constants for Snowpark Connect custom error codes."""
+
+    # 1000-1999: Startup related errors
+    MISSING_DATABASE = 1001
+    MISSING_SCHEMA = 1002
+    RESOURCE_INITIALIZATION_FAILED = 1003
+    TCP_PORT_ALREADY_IN_USE = 1004
+    INVALID_SPARK_CONNECT_URL = 1005
+    INVALID_STARTUP_INPUT = 1006
+    INVALID_STARTUP_OPERATION = 1007
+    STARTUP_CONNECTION_FAILED = 1008
+
+    # 2000-2999: Configuration related errors
+    INVALID_CONFIG_VALUE = 2001
+    CONFIG_CHANGE_NOT_ALLOWED = 2002
+    CONFIG_NOT_ENABLED = 2003
+
+    # 3000-3999: User code errors
+    INVALID_SQL_SYNTAX = 3001
+    TYPE_MISMATCH = 3002
+    INVALID_CAST = 3003
+    INVALID_FUNCTION_ARGUMENT = 3004
+    ARRAY_INDEX_OUT_OF_BOUNDS = 3005
+    DIVISION_BY_ZERO = 3006
+    INVALID_INPUT = 3007
+    INVALID_OPERATION = 3008
+    INSUFFICIENT_INPUT = 3009
+
+    # 4000-4999: What we don't support
+    UNSUPPORTED_OPERATION = 4001
+    UNSUPPORTED_TYPE = 4002
+
+    # 5000-5999: Internal errors
+    INTERNAL_ERROR = 5001
+    TABLE_NOT_FOUND = 5002
+    COLUMN_NOT_FOUND = 5003
+    AMBIGUOUS_COLUMN_NAME = 5004