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

snowflake/snowpark_connect/server.py

@@ -56,7 +56,11 @@ from snowflake.snowpark_connect.analyze_plan.map_tree_string import map_tree_str
 from snowflake.snowpark_connect.config import route_config_proto
 from snowflake.snowpark_connect.constants import SERVER_SIDE_SESSION_ID
 from snowflake.snowpark_connect.control_server import ControlServicer
-from snowflake.snowpark_connect.error.error_utils import build_grpc_error_response
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import (
+    attach_custom_error_code,
+    build_grpc_error_response,
+)
 from snowflake.snowpark_connect.execute_plan.map_execution_command import (
     map_execution_command,
 )
@@ -96,7 +100,7 @@ from snowflake.snowpark_connect.utils.interrupt import (
     interrupt_queries_with_tag,
     interrupt_query,
 )
-from snowflake.snowpark_connect.utils.profiling import profile_method
+from snowflake.snowpark_connect.utils.profiling import PROFILING_ENABLED, profile_method
 from snowflake.snowpark_connect.utils.session import (
     configure_snowpark_session,
     get_or_create_snowpark_session,
@@ -154,9 +158,9 @@ def _handle_exception(context, e: Exception):
         logger.error("Error: %s - %s", type(e).__name__, str(e))
 
     telemetry.report_request_failure(e)
-
     if tcm.TCM_MODE:
-        # TODO: SNOW-2009834 gracefully return error back in TCM
+        # spark decoder will catch the error and return it to GS gracefully
+        attach_custom_error_code(e, ErrorCodes.INTERNAL_ERROR)
         raise e
 
     from grpc_status import rpc_status
@@ -374,9 +378,13 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
                         ),
                     )
                 case _:
-                    raise SnowparkConnectNotImplementedError(
+                    exception = SnowparkConnectNotImplementedError(
                         f"ANALYZE PLAN NOT IMPLEMENTED:\n{request}"
                     )
+                    attach_custom_error_code(
+                        exception, ErrorCodes.UNSUPPORTED_OPERATION
+                    )
+                    raise exception
         except Exception as e:
             _handle_exception(context, e)
         finally:
@@ -527,9 +535,13 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
                             ),
                         )
                 case _:
-                    raise ValueError(
+                    exception = ValueError(
                         f"Unexpected payload type in AddArtifacts: {request.WhichOneof('payload')}"
                     )
+                    attach_custom_error_code(
+                        exception, ErrorCodes.UNSUPPORTED_OPERATION
+                    )
+                    raise exception
 
         for name, data in cache_data.items():
             _try_handle_local_relation(name, bytes(data))
@@ -635,9 +647,13 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
                 case proto_base.InterruptRequest.InterruptType.INTERRUPT_TYPE_OPERATION_ID:
                     interrupted_ids = interrupt_query(request.operation_id)
                 case _:
-                    raise SnowparkConnectNotImplementedError(
+                    exception = SnowparkConnectNotImplementedError(
                         f"INTERRUPT NOT IMPLEMENTED:\n{request}"
                     )
+                    attach_custom_error_code(
+                        exception, ErrorCodes.UNSUPPORTED_OPERATION
+                    )
+                    raise exception
 
             return proto_base.InterruptResponse(
                 session_id=request.session_id,
@@ -655,9 +671,11 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
         continue. If there is a ResultComplete, the client should use ReleaseExecute with
         """
         logger.info("ReattachExecute")
-        raise SnowparkConnectNotImplementedError(
+        exception = SnowparkConnectNotImplementedError(
             "Spark client has detached, please resubmit request. In a future version, the server will be support the reattach."
         )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     def ReleaseExecute(self, request: proto_base.ReleaseExecuteRequest, context):
         """Release an reattachable execution, or parts thereof.
@@ -760,8 +778,11 @@ def _serve(
 
         ChannelBuilder.MAX_MESSAGE_LENGTH = grpc_max_msg_size
 
+        # cProfile doesn't work correctly with multiple threads
+        max_workers = 1 if PROFILING_ENABLED else 10
+
         server = grpc.server(
-            futures.ThreadPoolExecutor(max_workers=10), options=server_options
+            futures.ThreadPoolExecutor(max_workers=max_workers), options=server_options
         )
         control_servicer = ControlServicer(session)
         proto_base_grpc.add_SparkConnectServiceServicer_to_server(
@@ -791,6 +812,7 @@ def _serve(
             )
         else:
             logger.error("Error starting up Snowpark Connect server", exc_info=True)
+        attach_custom_error_code(e, ErrorCodes.INTERNAL_ERROR)
         raise e
     finally:
         # flush the telemetry queue if possible
@@ -808,7 +830,9 @@ def _set_remote_url(remote_url: str):
     elif parsed_url.scheme == "unix":
         _server_url = remote_url.split("/;")[0]
     else:
-        raise RuntimeError(f"Invalid Snowpark Connect URL: {remote_url}")
+        exception = RuntimeError(f"Invalid Snowpark Connect URL: {remote_url}")
+        attach_custom_error_code(exception, ErrorCodes.INVALID_SPARK_CONNECT_URL)
+        raise exception
 
 
 def _set_server_tcp_port(server_port: int):
@@ -822,7 +846,9 @@ def _check_port_is_free(port: int) -> None:
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
         s.settimeout(1)
         if s.connect_ex(("127.0.0.1", port)) == 0:
-            raise RuntimeError(f"TCP port {port} is already in use")
+            exception = RuntimeError(f"TCP port {port} is already in use")
+            attach_custom_error_code(exception, ErrorCodes.TCP_PORT_ALREADY_IN_USE)
+            raise exception
 
 
 def _set_server_unix_domain_socket(path: str):
@@ -834,14 +860,18 @@ def _set_server_unix_domain_socket(path: str):
 def get_server_url() -> str:
     global _server_url
     if not _server_url:
-        raise RuntimeError("Server URL not set")
+        exception = RuntimeError("Server URL not set")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
     return _server_url
 
 
 def get_client_url() -> str:
     global _client_url
     if not _client_url:
-        raise RuntimeError("Client URL not set")
+        exception = RuntimeError("Client URL not set")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
     return _client_url
 
 
@@ -871,12 +901,14 @@ class UnixDomainSocketChannelBuilder(ChannelBuilder):
         if url is None:
             url = get_client_url()
         if url[:6] != "unix:/" or len(url) < 7:
-            raise PySparkValueError(
+            exception = PySparkValueError(
                 error_class="INVALID_CONNECT_URL",
                 message_parameters={
                     "detail": "The URL must start with 'unix://'. Please update the URL to follow the correct format, e.g., 'unix://unix_domain_socket_path'.",
                 },
             )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_SPARK_CONNECT_URL)
+            raise exception
 
         # Rewrite the URL to use http as the scheme so that we can leverage
         # Python's built-in parser to parse URL parameters
@@ -919,7 +951,7 @@ class UnixDomainSocketChannelBuilder(ChannelBuilder):
             for p in parts:
                 kv = p.split("=")
                 if len(kv) != 2:
-                    raise PySparkValueError(
+                    exception = PySparkValueError(
                         error_class="INVALID_CONNECT_URL",
                         message_parameters={
                             "detail": f"Parameter '{p}' should be provided as a "
@@ -927,6 +959,10 @@ class UnixDomainSocketChannelBuilder(ChannelBuilder):
                             f"the parameter to follow the correct format, e.g., 'key=value'.",
                         },
                     )
+                    attach_custom_error_code(
+                        exception, ErrorCodes.INVALID_SPARK_CONNECT_URL
+                    )
+                    raise exception
                 self.params[kv[0]] = urllib.parse.unquote(kv[1])
 
         netloc = self.url.netloc.split(":")
@@ -942,7 +978,7 @@ class UnixDomainSocketChannelBuilder(ChannelBuilder):
             self.host = netloc[0]
             self.port = int(netloc[1])
         else:
-            raise PySparkValueError(
+            exception = PySparkValueError(
                 error_class="INVALID_CONNECT_URL",
                 message_parameters={
                     "detail": f"Target destination '{self.url.netloc}' should match the "
@@ -950,6 +986,8 @@ class UnixDomainSocketChannelBuilder(ChannelBuilder):
                     f"the correct format, e.g., 'hostname:port'.",
                 },
             )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_SPARK_CONNECT_URL)
+            raise exception
 
     # We override this to enable compatibility with Spark 4.0
     host = None
@@ -988,9 +1026,11 @@ def start_jvm():
         if tcm.TCM_MODE:
             # No-op if JVM is already started in TCM mode
             return
-        raise RuntimeError(
+        exception = RuntimeError(
             "JVM must not be running when starting the Spark Connect server"
         )
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     pyspark_jars = (
         pathlib.Path(snowflake.snowpark_connect.__file__).parent / "includes/jars"
@@ -1065,14 +1105,18 @@ def start_session(
         _SPARK_CONNECT_GRPC_MAX_MESSAGE_SIZE = max_grpc_message_size
 
         if os.environ.get("SPARK_ENV_LOADED"):
-            raise RuntimeError(
+            exception = RuntimeError(
                 "Snowpark Connect cannot be run inside of a Spark environment"
             )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_STARTUP_OPERATION)
+            raise exception
         if connection_parameters is not None:
             if snowpark_session is not None:
-                raise ValueError(
+                exception = ValueError(
                     "Only specify one of snowpark_session and connection_parameters"
                 )
+                attach_custom_error_code(exception, ErrorCodes.INVALID_STARTUP_INPUT)
+                raise exception
             snowpark_session = snowpark.Session.builder.configs(
                 connection_parameters
             ).create()
@@ -1084,9 +1128,11 @@ def start_session(
             return
 
         if len(list(filter(None, [remote_url, tcp_port, unix_domain_socket]))) > 1:
-            raise RuntimeError(
+            exception = RuntimeError(
                 "Can only set at most one of remote_url, tcp_port, and unix_domain_socket"
             )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_STARTUP_INPUT)
+            raise exception
 
         url_from_env = os.environ.get("SPARK_REMOTE", None)
         if remote_url:
@@ -1124,7 +1170,11 @@ def start_session(
             server_thread.start()
             _server_running.wait()
             if _server_error:
-                raise RuntimeError("Snowpark Connect session failed to start")
+                exception = RuntimeError("Snowpark Connect session failed to start")
+                attach_custom_error_code(
+                    exception, ErrorCodes.STARTUP_CONNECTION_FAILED
+                )
+                raise exception
             return server_thread
         else:
             # Launch in the foreground.
@@ -1132,6 +1182,7 @@ def start_session(
     except Exception as e:
         _reset_server_run_state()
         logger.error(e, exc_info=True)
+        attach_custom_error_code(e, ErrorCodes.INTERNAL_ERROR)
         raise e
 
 
@@ -1165,6 +1216,7 @@ def get_session(url: Optional[str] = None, conf: SparkConf = None) -> SparkSessi
     except Exception as e:
         _reset_server_run_state()
         logger.error(e, exc_info=True)
+        attach_custom_error_code(e, ErrorCodes.INTERNAL_ERROR)
         raise e
 
 

snowflake/snowpark_connect/type_mapping.py

@@ -29,12 +29,17 @@ from snowflake.snowpark_connect.constants import (
 from snowflake.snowpark_connect.date_time_format_mapping import (
     convert_spark_format_to_snowflake,
 )
+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.literal import get_literal_field_and_name
 from snowflake.snowpark_connect.expression.map_sql_expression import (
     _INTERVAL_DAYTIME_PATTERN_RE,
     _INTERVAL_YEARMONTH_PATTERN_RE,
 )
-from snowflake.snowpark_connect.utils.context import get_is_evaluating_sql
+from snowflake.snowpark_connect.utils.context import (
+    get_is_evaluating_sql,
+    get_jpype_jclass_lock,
+)
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
@@ -61,12 +66,14 @@ SNOWPARK_TYPE_NAME_TO_PYSPARK_TYPE_NAME = {
 
 @cache
 def _get_struct_type_class():
-    return jpype.JClass("org.apache.spark.sql.types.StructType")
+    with get_jpype_jclass_lock():
+        return jpype.JClass("org.apache.spark.sql.types.StructType")
 
 
 @cache
 def get_python_sql_utils_class():
-    return jpype.JClass("org.apache.spark.sql.api.python.PythonSQLUtils")
+    with get_jpype_jclass_lock():
+        return jpype.JClass("org.apache.spark.sql.api.python.PythonSQLUtils")
 
 
 def _parse_ddl_with_spark_scala(ddl_string: str) -> pyspark.sql.types.DataType:
@@ -291,9 +298,11 @@ def snowpark_to_proto_type(
                 )
             }
         case _:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported snowpark data type: {data_type}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def cast_to_match_snowpark_type(
@@ -333,7 +342,9 @@ def cast_to_match_snowpark_type(
                 with suppress(TypeError):
                     date = datetime.strptime(content, format)
                     return date
-            raise ValueError(f"Date casting error for {str(content)}")
+            exception = ValueError(f"Date casting error for {str(content)}")
+            attach_custom_error_code(exception, ErrorCodes.INVALID_CAST)
+            raise exception
         case snowpark.types.ShortType:
             return int(content)
         case snowpark.types.StringType:
@@ -363,9 +374,11 @@ def cast_to_match_snowpark_type(
         case snowpark.types.DayTimeIntervalType:
             return str(content)
         case _:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported snowpark data type in casting: {data_type}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def snowpark_to_iceberg_type(data_type: snowpark.types.DataType) -> str:
@@ -398,9 +411,11 @@ def snowpark_to_iceberg_type(data_type: snowpark.types.DataType) -> str:
         case snowpark.types.TimestampType:
             return "timestamp"
         case _:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported snowpark data type for iceber: {data_type}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def proto_to_snowpark_type(
@@ -487,9 +502,11 @@ def map_snowpark_types_to_pyarrow_types(
                     )
                 )
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Unsupported arrow type {pa_type} for snowpark ArrayType."
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+                raise exception
         case snowpark.types.BinaryType:
             return pa.binary()
         case snowpark.types.BooleanType:
@@ -530,9 +547,11 @@ def map_snowpark_types_to_pyarrow_types(
                     ),
                 )
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Unsupported arrow type {pa_type} for snowpark MapType."
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+                raise exception
         case snowpark.types.NullType:
             return pa.string()
         case snowpark.types.ShortType:
@@ -557,15 +576,20 @@ def map_snowpark_types_to_pyarrow_types(
                     ]
                 )
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Unsupported arrow type {pa_type} for snowpark StructType."
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+                raise exception
         case snowpark.types.TimestampType:
-            unit = pa_type.unit
-            tz = pa_type.tz
+            # Check if pa_type has unit attribute (it should be a timestamp type)
+            unit = pa_type.unit if hasattr(pa_type, "unit") else "us"
+            tz = pa_type.tz if hasattr(pa_type, "tz") else None
+
+            # Spark truncates nanosecond precision to microseconds
             if unit == "ns":
-                # Spark truncates nanosecond precision to microseconds
                 unit = "us"
+
             return pa.timestamp(unit, tz=tz)
         case snowpark.types.VariantType:
             return pa.string()
@@ -576,9 +600,11 @@ def map_snowpark_types_to_pyarrow_types(
             # Return string type so formatted intervals are preserved in display
             return pa.string()
         case _:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported snowpark data type: {snowpark_type}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def map_pyarrow_to_snowpark_types(pa_type: pa.DataType) -> snowpark.types.DataType:
@@ -647,10 +673,15 @@ def map_pyarrow_to_snowpark_types(pa_type: pa.DataType) -> snowpark.types.DataTy
         return snowpark.types.TimestampType()
     elif pa.types.is_null(pa_type):
         return snowpark.types.NullType()
+    elif pa.types.is_duration(pa_type):
+        # Map PyArrow duration[us] to DayTimeIntervalType
+        return snowpark.types.DayTimeIntervalType()
     else:
-        raise SnowparkConnectNotImplementedError(
+        exception = SnowparkConnectNotImplementedError(
             f"Unsupported PyArrow data type: {pa_type}"
         )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
 
 def map_pyspark_types_to_snowpark_types(
@@ -736,9 +767,11 @@ def map_pyspark_types_to_snowpark_types(
         return snowpark.types.DayTimeIntervalType(
             type_to_map.startField, type_to_map.endField
         )
-    raise SnowparkConnectNotImplementedError(
+    exception = SnowparkConnectNotImplementedError(
         f"Unsupported spark data type: {type_to_map}"
     )
+    attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+    raise exception
 
 
 def map_snowpark_to_pyspark_types(
@@ -811,7 +844,11 @@ def map_snowpark_to_pyspark_types(
         return pyspark.sql.types.DayTimeIntervalType(
             type_to_map.start_field, type_to_map.end_field
         )
-    raise SnowparkConnectNotImplementedError(f"Unsupported data type: {type_to_map}")
+    exception = SnowparkConnectNotImplementedError(
+        f"Unsupported data type: {type_to_map}"
+    )
+    attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+    raise exception
 
 
 def map_simple_types(simple_type: str) -> snowpark.types.DataType:
@@ -861,14 +898,43 @@ def map_simple_types(simple_type: str) -> snowpark.types.DataType:
             return snowpark.types.YearMonthIntervalType()
         case type_name if _INTERVAL_DAYTIME_PATTERN_RE.match(type_name):
             return snowpark.types.DayTimeIntervalType()
+        # Year-Month interval cases
+        case "interval year":
+            return snowpark.types.YearMonthIntervalType(0)
+        case "interval month":
+            return snowpark.types.YearMonthIntervalType(1)
+        case "interval year to month":
+            return snowpark.types.YearMonthIntervalType(0, 1)
+        case "interval day":
+            return snowpark.types.DayTimeIntervalType(0)
+        case "interval hour":
+            return snowpark.types.DayTimeIntervalType(1)
+        case "interval minute":
+            return snowpark.types.DayTimeIntervalType(2)
+        case "interval second":
+            return snowpark.types.DayTimeIntervalType(3)
+        case "interval day to hour":
+            return snowpark.types.DayTimeIntervalType(0, 1)
+        case "interval day to minute":
+            return snowpark.types.DayTimeIntervalType(0, 2)
+        case "interval day to second":
+            return snowpark.types.DayTimeIntervalType(0, 3)
+        case "interval hour to minute":
+            return snowpark.types.DayTimeIntervalType(1, 2)
+        case "interval hour to second":
+            return snowpark.types.DayTimeIntervalType(1, 3)
+        case "interval minute to second":
+            return snowpark.types.DayTimeIntervalType(2, 3)
         case _:
             if simple_type.startswith("decimal"):
                 precision = int(simple_type.split("(")[1].split(",")[0])
                 scale = int(simple_type.split(",")[1].split(")")[0])
                 return snowpark.types.DecimalType(precision, scale)
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported simple type: {simple_type}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def map_json_schema_to_snowpark(
@@ -1009,9 +1075,11 @@ def map_spark_timestamp_format_expression(
             lit_value, _ = get_literal_field_and_name(arguments.literal)
             return convert_spark_format_to_snowflake(lit_value, timestamp_input_type)
         case other:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported expression type {other} in timestamp format argument"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def map_spark_number_format_expression(
@@ -1030,9 +1098,11 @@ def map_spark_number_format_expression(
             case "literal":
                 lit_value, _ = get_literal_field_and_name(arguments.literal)
             case other:
-                raise SnowparkConnectNotImplementedError(
+                exception = SnowparkConnectNotImplementedError(
                     f"Unsupported expression type {other} in number format argument"
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+                raise exception
 
     return _map_spark_to_snowflake_number_format(lit_value)
 

snowflake/snowpark_connect/typed_column.py

@@ -8,6 +8,7 @@ from functools import cached_property
 import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
 from snowflake.snowpark.column import Column
+from snowflake.snowpark_connect.column_qualifier import ColumnQualifier
 
 _EMPTY_COLUMN = Column("")
 
@@ -44,11 +45,11 @@ class TypedColumn:
     def alias(self, alias_name: str):
         return TypedColumn(self.col.alias(alias_name), self._type_resolver)
 
-    def set_qualifiers(self, qualifiers: list[str]):
+    def set_qualifiers(self, qualifiers: set[ColumnQualifier]) -> None:
         self.qualifiers = qualifiers
 
-    def get_qualifiers(self) -> list[str]:
-        return getattr(self, "qualifiers", [])
+    def get_qualifiers(self) -> set[ColumnQualifier]:
+        return getattr(self, "qualifiers", {ColumnQualifier.no_qualifier()})
 
     def set_catalog_database_info(self, catalog_database_info: dict[str, str]) -> None:
         self._catalog_database_info = catalog_database_info
@@ -63,12 +64,13 @@ class TypedColumn:
     def get_database(self) -> str | None:
         return self._catalog_database_info.get("database")
 
-    def set_multi_col_qualifiers(self, qualifiers: list[list[str]]):
+    def set_multi_col_qualifiers(self, qualifiers: list[set[ColumnQualifier]]) -> None:
         self.multi_col_qualifiers = qualifiers
 
-    def get_multi_col_qualifiers(self, num_columns) -> list[list[str]]:
+    def get_multi_col_qualifiers(self, num_columns) -> list[set[ColumnQualifier]]:
         if not hasattr(self, "multi_col_qualifiers"):
-            return [[] for i in range(num_columns)]
+
+            return [{ColumnQualifier.no_qualifier()} for i in range(num_columns)]
         assert (
             len(self.multi_col_qualifiers) == num_columns
         ), f"Expected {num_columns} multi-column qualifiers, got {len(self.multi_col_qualifiers)}"

snowflake/snowpark_connect/utils/context.py

@@ -2,10 +2,12 @@
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
 
+import os
 import re
+import threading
 from contextlib import contextmanager
 from contextvars import ContextVar
-from typing import Mapping, Optional
+from typing import Iterator, Mapping, Optional
 
 import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 
@@ -25,6 +27,7 @@ _is_evaluating_sql = ContextVar[bool]("_is_evaluating_sql", default=False)
 _is_evaluating_join_condition = ContextVar(
     "_is_evaluating_join_condition", default=("default", False, [], [])
 )
+_is_processing_order_by = ContextVar[bool]("_is_processing_order_by", default=False)
 
 _sql_aggregate_function_count = ContextVar[int](
     "_contains_aggregate_function", default=0
@@ -56,6 +59,23 @@ _is_in_pivot = ContextVar[bool]("_is_in_pivot", default=False)
 _is_in_udtf_context = ContextVar[bool]("_is_in_udtf_context", default=False)
 _accessing_temp_object = ContextVar[bool]("_accessing_temp_object", default=False)
 
+# Thread-safe lock for JPype JClass creation to prevent access violations
+_jpype_jclass_lock = threading.Lock()
+
+
+@contextmanager
+def get_jpype_jclass_lock() -> Iterator[None]:
+    """
+    Context manager that acquires the JPype JClass lock on Windows platforms.
+    On non-Windows (os.name != 'nt'), it yields without acquiring the lock.
+    """
+    if os.name == "nt":
+        with _jpype_jclass_lock:
+            yield
+    else:
+        yield
+
+
 # Lateral Column Alias helpers
 # We keep a thread-local mapping from alias name -> TypedColumn that is
 # populated incrementally while the projection list is being processed.
@@ -207,6 +227,27 @@ def push_evaluating_sql_scope():
         _is_evaluating_sql.set(prev)
 
 
+def get_is_processing_order_by() -> bool:
+    """
+    Gets the value of _is_processing_order_by for the current context, defaults to False.
+    """
+    return _is_processing_order_by.get()
+
+
+@contextmanager
+def push_processing_order_by_scope():
+    """
+    Context manager that sets a flag indicating if ORDER BY expressions are being evaluated.
+    This enables optimizations like reusing already-computed UDF columns.
+    """
+    prev = _is_processing_order_by.get()
+    try:
+        _is_processing_order_by.set(True)
+        yield
+    finally:
+        _is_processing_order_by.set(prev)
+
+
 def get_is_evaluating_join_condition() -> tuple[str, bool, list, list]:
     """
     Gets the value of _is_evaluating_join_condition for the current context, defaults to False.

snowflake/snowpark_connect/utils/describe_query_cache.py

@@ -12,6 +12,8 @@ from typing import Any
 from snowflake import snowpark
 from snowflake.connector.cursor import ResultMetadataV2
 from snowflake.snowpark._internal.server_connection import ServerConnection
+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.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.utils.telemetry import telemetry
@@ -148,6 +150,7 @@ def instrument_session_for_describe_cache(session: snowpark.Session):
                 telemetry.report_query(result, **kwargs)
             except Exception as e:
                 telemetry.report_query(e, **kwargs)
+                attach_custom_error_code(e, ErrorCodes.INTERNAL_ERROR)
                 raise e
             return result