snowpark-connect 0.21.0__py3-none-any.whl → 0.22.1__py3-none-any.whl

snowflake/snowpark_connect/config.py

@@ -168,6 +168,9 @@ class GlobalConfig:
         "snowpark.connect.udf.packages": lambda session, packages: session.add_packages(
             *packages.strip("[] ").split(",")
         ),
+        "snowpark.connect.udf.imports": lambda session, imports: parse_imports(
+            session, imports
+        ),
     }
 
     float_config_list = []
@@ -332,7 +335,7 @@ def route_config_proto(
     match op_type:
         case "set":
             logger.info("SET")
-
+            telemetry.report_config_set(config.operation.set.pairs)
             for pair in config.operation.set.pairs:
                 # Check if the value field is present, not present when invalid fields are set in conf.
                 if not pair.HasField("value"):
@@ -342,7 +345,6 @@ def route_config_proto(
                         f"Cannot set config '{pair.key}' to None"
                     )
 
-                telemetry.report_config_set(pair.key, pair.value)
                 set_config_param(
                     config.session_id, pair.key, pair.value, snowpark_session
                 )
@@ -350,14 +352,15 @@ def route_config_proto(
             return proto_base.ConfigResponse(session_id=config.session_id)
         case "unset":
             logger.info("UNSET")
+            telemetry.report_config_unset(config.operation.unset.keys)
             for key in config.operation.unset.keys:
-                telemetry.report_config_unset(key)
                 unset_config_param(config.session_id, key, snowpark_session)
 
             return proto_base.ConfigResponse(session_id=config.session_id)
         case "get":
             logger.info("GET")
             res = proto_base.ConfigResponse(session_id=config.session_id)
+            telemetry.report_config_get(config.operation.get.keys)
             for key in config.operation.get.keys:
                 pair = res.pairs.add()
                 pair.key = key
@@ -367,6 +370,9 @@ def route_config_proto(
             return res
         case "get_with_default":
             logger.info("GET_WITH_DEFAULT")
+            telemetry.report_config_get(
+                [pair.key for pair in config.operation.get_with_default.pairs]
+            )
             result_pairs = [
                 proto_base.KeyValue(
                     key=pair.key,
@@ -383,6 +389,7 @@ def route_config_proto(
         case "get_option":
             logger.info("GET_OPTION")
             res = proto_base.ConfigResponse(session_id=config.session_id)
+            telemetry.report_config_get(config.operation.get_option.keys)
             for key in config.operation.get_option.keys:
                 pair = res.pairs.add()
                 pair.key = key
@@ -411,6 +418,7 @@ def route_config_proto(
         case "is_modifiable":
             logger.info("IS_MODIFIABLE")
             res = proto_base.ConfigResponse(session_id=config.session_id)
+            telemetry.report_config_get(config.operation.is_modifiable.keys)
             for key in config.operation.is_modifiable.keys:
                 pair = res.pairs.add()
                 pair.key = key
@@ -588,3 +596,11 @@ def auto_uppercase_non_column_identifiers() -> bool:
     return session_config[
         "snowpark.connect.sql.identifiers.auto-uppercase"
     ].lower() in ("all", "all_except_columns")
+
+
+def parse_imports(session: snowpark.Session, imports: str | None) -> None:
+    if not imports:
+        return
+
+    for udf_import in imports.strip("[] ").split(","):
+        session.add_import(udf_import)

snowflake/snowpark_connect/error/error_utils.py

@@ -28,7 +28,9 @@ from pyspark.errors.exceptions.base import (
     PySparkException,
     PythonException,
     SparkRuntimeException,
+    UnsupportedOperationException,
 )
+from pyspark.errors.exceptions.connect import SparkConnectGrpcException
 from snowflake.core.exceptions import NotFoundError
 
 from snowflake.connector.errors import ProgrammingError
@@ -49,7 +51,9 @@ SPARK_PYTHON_TO_JAVA_EXCEPTION = {
     ArrayIndexOutOfBoundsException: "java.lang.ArrayIndexOutOfBoundsException",
     NumberFormatException: "java.lang.NumberFormatException",
     SparkRuntimeException: "org.apache.spark.SparkRuntimeException",
+    SparkConnectGrpcException: "pyspark.errors.exceptions.connect.SparkConnectGrpcException",
     PythonException: "org.apache.spark.api.python.PythonException",
+    UnsupportedOperationException: "java.lang.UnsupportedOperationException",
 }
 
 WINDOW_FUNCTION_ANALYSIS_EXCEPTION_SQL_ERROR_CODE = {1005, 2303}
@@ -68,6 +72,9 @@ init_multi_args_exception_pattern = (
 terminate_multi_args_exception_pattern = (
     r"terminate\(\) missing \d+ required positional argument"
 )
+snowpark_connect_exception_pattern = re.compile(
+    r"\[snowpark-connect-exception(?::(\w+))?\]\s*(.+?)'\s*is not recognized"
+)
 
 
 def contains_udtf_select(sql_string):
@@ -100,6 +107,19 @@ def _get_converted_known_sql_or_custom_exception(
         return SparkRuntimeException(
             message="Unexpected value for start in function slice: SQL array indices start at 1."
         )
+    match = snowpark_connect_exception_pattern.search(
+        ex.message if hasattr(ex, "message") else str(ex)
+    )
+    if match:
+        class_name = match.group(1)
+        message = match.group(2)
+        exception_class = (
+            globals().get(class_name, SparkConnectGrpcException)
+            if class_name
+            else SparkConnectGrpcException
+        )
+        return exception_class(message=message)
+
     if "select with no columns" in msg and contains_udtf_select(query):
         # We try our best to detect if the SQL string contains a UDTF call and the output schema is empty.
         return PythonException(message=f"[UDTF_RETURN_SCHEMA_MISMATCH] {ex.message}")
@@ -131,6 +151,11 @@ def _get_converted_known_sql_or_custom_exception(
                 message=f"[UDTF_EXEC_ERROR] User defined table function encountered an error in the terminate method: {ex.message}"
             )
 
+        if "failed to split string, provided pattern:" in msg:
+            return IllegalArgumentException(
+                message=f"Failed to split string using provided pattern. {ex.message}"
+            )
+
         if "100357" in msg and "wrong tuple size for returned value" in msg:
             return PythonException(
                 message=f"[UDTF_RETURN_SCHEMA_MISMATCH] The number of columns in the result does not match the specified schema. {ex.message}"

snowflake/snowpark_connect/expression/map_udf.py

@@ -13,10 +13,7 @@ from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.expression.typer import ExpressionTyper
 from snowflake.snowpark_connect.type_mapping import proto_to_snowpark_type
 from snowflake.snowpark_connect.typed_column import TypedColumn
-from snowflake.snowpark_connect.utils.session import (
-    get_or_create_snowpark_session,
-    get_python_udxf_import_files,
-)
+from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
 from snowflake.snowpark_connect.utils.udf_helper import (
     SnowparkUDF,
     gen_input_types,
@@ -28,6 +25,9 @@ from snowflake.snowpark_connect.utils.udf_helper import (
 from snowflake.snowpark_connect.utils.udf_utils import (
     ProcessCommonInlineUserDefinedFunction,
 )
+from snowflake.snowpark_connect.utils.udxf_import_utils import (
+    get_python_udxf_import_files,
+)
 
 
 def process_udf_return_type(