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

snowflake/snowpark_connect/utils/scala_udf_utils.py

@@ -15,14 +15,21 @@ Key components:
 - Type mapping functions for different type systems
 - UDF creation and management utilities
 """
-import re
 from dataclasses import dataclass
-from enum import Enum
 from typing import List, Union
 
 import snowflake.snowpark.types as snowpark_type
 import snowflake.snowpark_connect.includes.python.pyspark.sql.connect.proto.types_pb2 as types_proto
-from snowflake.snowpark_connect.resources_initializer import RESOURCE_PATH
+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 map_type_to_snowflake_type
+from snowflake.snowpark_connect.utils.jvm_udf_utils import (
+    NullHandling,
+    Param,
+    ReturnType,
+    Signature,
+    build_jvm_udxf_imports,
+)
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.utils.udf_utils import (
     ProcessCommonInlineUserDefinedFunction,
@@ -59,58 +66,6 @@ class ScalaUdf:
         self._return_type = return_type
 
 
-@dataclass(frozen=True)
-class Param:
-    """
-    Represents a function parameter with name and data type.
-
-    Attributes:
-        name: Parameter name
-        data_type: Parameter data type as a string
-    """
-
-    name: str
-    data_type: str
-
-
-@dataclass(frozen=True)
-class NullHandling(str, Enum):
-    """
-    Enumeration for UDF null handling behavior.
-
-    Determines how the UDF behaves when input parameters contain null values.
-    """
-
-    RETURNS_NULL_ON_NULL_INPUT = "RETURNS NULL ON NULL INPUT"
-    CALLED_ON_NULL_INPUT = "CALLED ON NULL INPUT"
-
-
-@dataclass(frozen=True)
-class ReturnType:
-    """
-    Represents the return type of a function.
-
-    Attributes:
-        data_type: Return data type as a string
-    """
-
-    data_type: str
-
-
-@dataclass(frozen=True)
-class Signature:
-    """
-    Represents a function signature with parameters and return type.
-
-    Attributes:
-        params: List of function parameters
-        returns: Function return type
-    """
-
-    params: List[Param]
-    returns: ReturnType
-
-
 @dataclass(frozen=True)
 class ScalaUDFDef:
     """
@@ -147,76 +102,45 @@ class ScalaUDFDef:
             String containing the complete Scala code for the UDF body
         """
         # Convert Array to Seq for Scala compatibility in function signatures.
-        udf_func_input_types = (
-            ", ".join(p.data_type for p in self.scala_signature.params)
-        ).replace("Array", "Seq")
+        # Replace each "Variant" type with "Any" in the function signature since fromVariant returns Any
+        udf_func_input_types = ", ".join(
+            "Any" if p.data_type == "Variant" else p.data_type.replace("Array", "Seq")
+            for p in self.scala_signature.params
+        )
+        udf_func_return_type = self.scala_signature.returns.data_type.replace(
+            "Array", "Seq"
+        )
+
         # Create the Scala arguments and input types string: "arg0: Type0, arg1: Type1, ...".
         joined_wrapper_arg_and_input_types_str = ", ".join(
             f"{p.name}: {p.data_type}" for p in self.scala_signature.params
         )
-        # This is used in defining the input types for the wrapper function. For Maps to work correctly with Scala UDFs,
-        # we need to set the Map types to Map[String, String]. These get cast to the respective original types
-        # when the original UDF function is invoked.
-        wrapper_arg_and_input_types_str = re.sub(
-            pattern=r"Map\[\w+,\s\w+\]",
-            repl="Map[String, String]",
-            string=joined_wrapper_arg_and_input_types_str,
-        )
-        invocation_args = ", ".join(self.scala_invocation_args)
-
-        # Cannot directly return a map from a Scala UDF due to issues with non-String values. Snowflake SQL Scala only
-        # supports Map[String, String] as input types. Therefore, we convert the map to a JSON string before returning.
-        # This is processed as a Variant by SQL.
-        udf_func_return_type = self.scala_signature.returns.data_type
-        is_map_return = udf_func_return_type.startswith("Map")
-        wrapper_return_type = "String" if is_map_return else udf_func_return_type
-
-        # Need to call the map to JSON string converter when a map is returned by the user's function.
-        invoke_udf_func = (
-            f"write(func({invocation_args}))"
-            if is_map_return
-            else f"func({invocation_args})"
-        )
 
-        # The lines of code below are required only when a Map is returned by the UDF. This is needed to serialize the
-        # map output to a JSON string.
-        map_return_imports = (
-            ""
-            if not is_map_return
-            else """
-import org.json4s._
-import org.json4s.native.Serialization._
-import org.json4s.native.Serialization
-"""
-        )
-        map_return_formatter = (
-            ""
-            if not is_map_return
-            else """
-  implicit val formats = Serialization.formats(NoTypeHints)
-"""
-        )
+        # All Scala UDFs return Variant to ensure consistency and avoid type conversion issues.
+        wrapper_return_type = "Variant"
+        wrapped_args = [
+            f"udfPacket.fromVariant({arg}, {i})" if p.data_type == "Variant" else arg
+            for i, (arg, p) in enumerate(
+                zip(self.scala_invocation_args, self.scala_signature.params)
+            )
+        ]
+        invocation_args = ", ".join(wrapped_args)
+        invoke_udf_func = f"func({invocation_args})"
+
+        # Always wrap the result in Utils.toVariant() to ensure all Scala UDFs return Variant
+        invoke_udf_func = f"Utils.toVariant({invoke_udf_func})"
 
-        return f"""import org.apache.spark.sql.connect.common.UdfPacket
-{map_return_imports}
-import java.io.{{ByteArrayInputStream, ObjectInputStream}}
-import java.nio.file.{{Files, Paths}}
+        return f"""
+import org.apache.spark.sql.connect.common.UdfPacket
+import com.snowflake.sas.scala.UdfPacketUtils._
+import com.snowflake.sas.scala.Utils
+import com.snowflake.snowpark_java.types.Variant
 
 object __RecreatedSparkUdf {{
-  {map_return_formatter}
-  private lazy val func: ({udf_func_input_types}) => {udf_func_return_type} = {{
-    val importDirectory = System.getProperty("com.snowflake.import_directory")
-    val fPath = importDirectory + "{self.name}.bin"
-    val bytes = Files.readAllBytes(Paths.get(fPath))
-    val ois = new ObjectInputStream(new ByteArrayInputStream(bytes))
-    try {{
-      ois.readObject().asInstanceOf[UdfPacket].function.asInstanceOf[({udf_func_input_types}) => {udf_func_return_type}]
-    }} finally {{
-      ois.close()
-    }}
-  }}
+  private lazy val udfPacket: UdfPacket = Utils.deserializeUdfPacket("{self.name}.bin")
+  private lazy val func: ({udf_func_input_types}) => {udf_func_return_type} = udfPacket.function.asInstanceOf[({udf_func_input_types}) => {udf_func_return_type}]
 
-  def __wrapperFunc({wrapper_arg_and_input_types_str}): {wrapper_return_type} = {{
+  def __wrapperFunc({joined_wrapper_arg_and_input_types_str}): {wrapper_return_type} = {{
     {invoke_udf_func}
   }}
 }}
@@ -260,70 +184,6 @@ $$
 $$;"""
 
 
-def build_scala_udf_imports(session, payload, udf_name, is_map_return) -> List[str]:
-    """
-    Build the list of imports needed for the Scala UDF.
-
-    This function:
-    1. Saves the UDF payload to a binary file in the session stage
-    2. Collects user-uploaded JAR files from the stage
-    3. Returns a list of all required JAR files for the UDF
-
-    Args:
-        session: Snowpark session
-        payload: Binary payload containing the serialized UDF
-        udf_name: Name of the UDF (used for the binary file name)
-        is_map_return: Indicates if the UDF returns a Map (affects imports)
-
-    Returns:
-        List of JAR file paths to be imported by the UDF
-    """
-    # Save pciudf._payload to a bin file:
-    import io
-
-    payload_as_stream = io.BytesIO(payload)
-    stage = session.get_session_stage()
-    stage_resource_path = stage + RESOURCE_PATH
-    closure_binary_file = stage_resource_path + "/" + udf_name + ".bin"
-    session.file.put_stream(
-        payload_as_stream,
-        closure_binary_file,
-        overwrite=True,
-    )
-
-    # Get a list of the jar files uploaded to the stage. We need to import the user's jar for the Scala UDF.
-    res = session.sql(rf"LIST {stage}/ PATTERN='.*\.jar';").collect()
-    user_jars = []
-    for row in res:
-        if RESOURCE_PATH not in row[0]:
-            # Remove the stage path since it is not properly formatted.
-            user_jars.append(row[0][row[0].find("/") :])
-
-    # Jars used when the return type is a Map.
-    map_jars = (
-        []
-        if not is_map_return
-        else [
-            f"{stage_resource_path}/json4s-core_2.12-3.7.0-M11.jar",
-            f"{stage_resource_path}/json4s-native_2.12-3.7.0-M11.jar",
-            f"{stage_resource_path}/paranamer-2.8.3.jar",
-        ]
-    )
-
-    # Format the user jars to be used in the IMPORTS clause of the stored procedure.
-    return (
-        [
-            closure_binary_file,
-            f"{stage_resource_path}/spark-connect-client-jvm_2.12-3.5.6.jar",
-            f"{stage_resource_path}/spark-common-utils_2.12-3.5.6.jar",
-            f"{stage_resource_path}/spark-sql_2.12-3.5.6.jar",
-            f"{stage_resource_path}/json4s-ast_2.12-3.7.0-M11.jar",
-        ]
-        + map_jars
-        + [f"{stage + jar}" for jar in user_jars]
-    )
-
-
 def create_scala_udf(pciudf: ProcessCommonInlineUserDefinedFunction) -> ScalaUdf:
     """
     Create a Scala UDF in Snowflake from a ProcessCommonInlineUserDefinedFunction object.
@@ -343,7 +203,13 @@ def create_scala_udf(pciudf: ProcessCommonInlineUserDefinedFunction) -> ScalaUdf
     Returns:
         A ScalaUdf object representing the created or cached Scala UDF.
     """
-    from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
+    from snowflake.snowpark_connect.resources_initializer import (
+        ensure_scala_udf_jars_uploaded,
+    )
+
+    # Lazily upload Scala UDF jars on-demand when a Scala UDF is actually created.
+    # This is thread-safe and will only upload once even if multiple threads call it.
+    ensure_scala_udf_jars_uploaded()
 
     function_name = pciudf._function_name
     # If a function name is not provided, hash the binary file and use the first ten characters as the function name.
@@ -353,11 +219,6 @@ def create_scala_udf(pciudf: ProcessCommonInlineUserDefinedFunction) -> ScalaUdf
         function_name = hashlib.sha256(pciudf._payload).hexdigest()[:10]
     udf_name = CREATE_SCALA_UDF_PREFIX + function_name
 
-    session = get_or_create_snowpark_session()
-    if udf_name in session._udfs:
-        cached_udf = session._udfs[udf_name]
-        return ScalaUdf(cached_udf.name, cached_udf.input_types, cached_udf.return_type)
-
     # In case the Scala UDF was created with `spark.udf.register`, the Spark Scala input types (from protobuf) are
     # stored in pciudf.scala_input_types.
     # We cannot rely solely on the inputTypes field from the Scala UDF or the Snowpark input types, since:
@@ -376,30 +237,40 @@ def create_scala_udf(pciudf: ProcessCommonInlineUserDefinedFunction) -> ScalaUdf
             param_name = "arg" + str(i)
             # Create the Scala arguments and input types string: "arg0: Type0, arg1: Type1, ...".
             scala_input_params.append(
-                Param(param_name, map_type_to_scala_type(input_type))
+                Param(param_name, _map_type_to_scala_type(input_type, is_input=True))
             )
             # Create the Snowflake SQL arguments and input types string: "arg0 TYPE0, arg1 TYPE1, ...".
-            sql_input_params.append(
-                Param(param_name, map_type_to_snowflake_type(input_type))
+            # For arrays and structs, use VARIANT type in SQL signature
+            is_snowpark_type = isinstance(input_type, snowpark_type.DataType)
+            is_array = (
+                is_snowpark_type and isinstance(input_type, snowpark_type.ArrayType)
+            ) or (not is_snowpark_type and input_type.WhichOneof("kind") == "array")
+            is_map = (
+                is_snowpark_type and isinstance(input_type, snowpark_type.MapType)
+            ) or (not is_snowpark_type and input_type.WhichOneof("kind") == "map")
+            sql_type = (
+                "VARIANT"
+                if is_array or is_map
+                else map_type_to_snowflake_type(input_type)
             )
+            sql_input_params.append(Param(param_name, sql_type))
             # In the case of Map input types, we need to cast the argument to the correct type in Scala.
-            # Snowflake SQL Scala can only handle MAP[VARCHAR, VARCHAR] as input types.
-            scala_invocation_args.append(
-                cast_scala_map_args_from_given_type(param_name, input_type)
-            )
+            scala_invocation_args.append(param_name)
+
+    scala_return_type = _map_type_to_scala_type(
+        pciudf._original_return_type, is_input=False
+    )
+    # All Scala UDFs now return VARIANT to ensure consistency and avoid type conversion issues.
+    # The actual type conversion is handled after the UDF is called.
+    from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
 
-    scala_return_type = map_type_to_scala_type(pciudf._original_return_type)
-    # If the SQL return type is a MAP, change this to VARIANT because of issues with Scala UDFs.
-    sql_return_type = map_type_to_snowflake_type(pciudf._original_return_type)
-    imports = build_scala_udf_imports(
+    session = get_or_create_snowpark_session()
+    imports = build_jvm_udxf_imports(
         session,
         pciudf._payload,
         udf_name,
-        is_map_return=sql_return_type.startswith("MAP"),
-    )
-    sql_return_type = (
-        "VARIANT" if sql_return_type.startswith("MAP") else sql_return_type
     )
+    sql_return_type = "VARIANT"
 
     udf_def = ScalaUDFDef(
         name=udf_name,
@@ -418,21 +289,30 @@ def create_scala_udf(pciudf: ProcessCommonInlineUserDefinedFunction) -> ScalaUdf
     return ScalaUdf(udf_name, pciudf._input_types, pciudf._return_type)
 
 
-def map_type_to_scala_type(
-    t: Union[snowpark_type.DataType, types_proto.DataType]
+def _map_type_to_scala_type(
+    t: Union[snowpark_type.DataType, types_proto.DataType], is_input: bool = False
 ) -> str:
-    """Maps a Snowpark or Spark protobuf type to a Scala type string."""
+    """Maps a Snowpark or Spark protobuf type to a Scala type string.
+
+    Args:
+        t: The type to map
+        is_input: If True, maps array types to Variant (for UDF inputs).
+                  If False, maps array types to Array[ElementType] (for UDF outputs).
+    """
     if not t:
         return "String"
     is_snowpark_type = isinstance(t, snowpark_type.DataType)
     condition = type(t) if is_snowpark_type else t.WhichOneof("kind")
     match condition:
         case snowpark_type.ArrayType | "array":
-            return (
-                f"Array[{map_type_to_scala_type(t.element_type)}]"
-                if is_snowpark_type
-                else f"Array[{map_type_to_scala_type(t.array.element_type)}]"
-            )
+            if is_input:
+                return "Variant"
+            else:
+                return (
+                    f"Array[{_map_type_to_scala_type(t.element_type, is_input=False)}]"
+                    if is_snowpark_type
+                    else f"Array[{_map_type_to_scala_type(t.array.element_type, is_input=False)}]"
+                )
         case snowpark_type.BinaryType | "binary":
             return "Array[Byte]"
         case snowpark_type.BooleanType | "boolean":
@@ -453,16 +333,18 @@ def map_type_to_scala_type(
             return "Int"
         case snowpark_type.LongType | "long":
             return "Long"
-        case snowpark_type.MapType | "map":  # can also map to OBJECT in Snowflake
+        case snowpark_type.MapType | "map":
+            if is_input:
+                return "Variant"
             key_type = (
-                map_type_to_scala_type(t.key_type)
+                _map_type_to_scala_type(t.key_type)
                 if is_snowpark_type
-                else map_type_to_scala_type(t.map.key_type)
+                else _map_type_to_scala_type(t.map.key_type)
             )
             value_type = (
-                map_type_to_scala_type(t.value_type)
+                _map_type_to_scala_type(t.value_type)
                 if is_snowpark_type
-                else map_type_to_scala_type(t.map.value_type)
+                else _map_type_to_scala_type(t.map.value_type)
             )
             return f"Map[{key_type}, {value_type}]"
         case snowpark_type.NullType | "null":
@@ -471,126 +353,13 @@ def map_type_to_scala_type(
             return "Short"
         case snowpark_type.StringType | "string" | "char" | "varchar":
             return "String"
+        case snowpark_type.StructType | "struct":
+            return "Variant"
         case snowpark_type.TimestampType | "timestamp" | "timestamp_ntz":
             return "java.sql.Timestamp"
         case snowpark_type.VariantType:
             return "Variant"
         case _:
-            raise ValueError(f"Unsupported Snowpark type: {t}")
-
-
-def map_type_to_snowflake_type(
-    t: Union[snowpark_type.DataType, types_proto.DataType]
-) -> str:
-    """Maps a Snowpark or Spark protobuf type to a Snowflake type string."""
-    if not t:
-        return "VARCHAR"
-    is_snowpark_type = isinstance(t, snowpark_type.DataType)
-    condition = type(t) if is_snowpark_type else t.WhichOneof("kind")
-    match condition:
-        case snowpark_type.ArrayType | "array":
-            return (
-                f"ARRAY({map_type_to_snowflake_type(t.element_type)})"
-                if is_snowpark_type
-                else f"ARRAY({map_type_to_snowflake_type(t.array.element_type)})"
-            )
-        case snowpark_type.BinaryType | "binary":
-            return "BINARY"
-        case snowpark_type.BooleanType | "boolean":
-            return "BOOLEAN"
-        case snowpark_type.ByteType | "byte":
-            return "TINYINT"
-        case snowpark_type.DateType | "date":
-            return "DATE"
-        case snowpark_type.DecimalType | "decimal":
-            return "NUMBER"
-        case snowpark_type.DoubleType | "double":
-            return "DOUBLE"
-        case snowpark_type.FloatType | "float":
-            return "FLOAT"
-        case snowpark_type.GeographyType:
-            return "GEOGRAPHY"
-        case snowpark_type.IntegerType | "integer":
-            return "INT"
-        case snowpark_type.LongType | "long":
-            return "BIGINT"
-        case snowpark_type.MapType | "map":
-            # Maps to OBJECT in Snowflake if key and value types are not specified.
-            key_type = (
-                map_type_to_snowflake_type(t.key_type)
-                if is_snowpark_type
-                else map_type_to_snowflake_type(t.map.key_type)
-            )
-            value_type = (
-                map_type_to_snowflake_type(t.value_type)
-                if is_snowpark_type
-                else map_type_to_snowflake_type(t.map.value_type)
-            )
-            return (
-                f"MAP({key_type}, {value_type})"
-                if key_type and value_type
-                else "OBJECT"
-            )
-        case snowpark_type.NullType | "null":
-            return "VARCHAR"
-        case snowpark_type.ShortType | "short":
-            return "SMALLINT"
-        case snowpark_type.StringType | "string" | "char" | "varchar":
-            return "VARCHAR"
-        case snowpark_type.TimestampType | "timestamp" | "timestamp_ntz":
-            return "TIMESTAMP"
-        case snowpark_type.VariantType:
-            return "VARIANT"
-        case _:
-            raise ValueError(f"Unsupported Snowpark type: {t}")
-
-
-def cast_scala_map_args_from_given_type(
-    arg_name: str, input_type: Union[snowpark_type.DataType, types_proto.DataType]
-) -> str:
-    """If the input_type is a Map, cast the argument arg_name to a Map[key_type, value_type] in Scala."""
-    is_snowpark_type = isinstance(input_type, snowpark_type.DataType)
-
-    def convert_from_string_to_type(
-        arg_name: str, t: Union[snowpark_type.DataType, types_proto.DataType]
-    ) -> str:
-        """Convert the string argument arg_name to the specified type t in Scala."""
-        condition = type(t) if is_snowpark_type else t.WhichOneof("kind")
-        match condition:
-            case snowpark_type.BinaryType | "binary":
-                return arg_name + ".getBytes()"
-            case snowpark_type.BooleanType | "boolean":
-                return arg_name + ".toBoolean"
-            case snowpark_type.ByteType | "byte":
-                return arg_name + ".getBytes().head"  # TODO: verify if this is correct
-            case snowpark_type.DateType | "date":
-                return f"java.sql.Date.valueOf({arg_name})"
-            case snowpark_type.DecimalType | "decimal":
-                return f"new BigDecimal({arg_name})"
-            case snowpark_type.DoubleType | "double":
-                return arg_name + ".toDouble"
-            case snowpark_type.FloatType | "float":
-                return arg_name + ".toFloat"
-            case snowpark_type.IntegerType | "integer":
-                return arg_name + ".toInt"
-            case snowpark_type.LongType | "long":
-                return arg_name + ".toLong"
-            case snowpark_type.ShortType | "short":
-                return arg_name + ".toShort"
-            case snowpark_type.StringType | "string" | "char" | "varchar":
-                return arg_name
-            case snowpark_type.TimestampType | "timestamp" | "timestamp_ntz":
-                return "java.sql.Timestamp.valueOf({arg_name})"
-            case _:
-                raise ValueError(f"Unsupported Snowpark type: {t}")
-
-    if (is_snowpark_type and isinstance(input_type, snowpark_type.MapType)) or (
-        not is_snowpark_type and input_type.WhichOneof("kind") == "map"
-    ):
-        key_type = input_type.key_type if is_snowpark_type else input_type.map.key_type
-        value_type = (
-            input_type.value_type if is_snowpark_type else input_type.map.value_type
-        )
-        return f"{arg_name}.map {{ case (k, v) => ({convert_from_string_to_type('k', key_type)}, {convert_from_string_to_type('v', value_type)})}}"
-    else:
-        return arg_name
+            exception = ValueError(f"Unsupported Snowpark type: {t}")
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+            raise exception

snowflake/snowpark_connect/utils/sequence.py

@@ -0,0 +1,21 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import threading
+from collections import defaultdict
+
+from snowflake.snowpark_connect.utils.context import get_spark_session_id
+
+# per session number sequences to generate unique snowpark columns
+_session_sequences = defaultdict(int)
+
+_lock = threading.Lock()
+
+
+def next_unique_num():
+    session_id = get_spark_session_id()
+    with _lock:
+        next_num = _session_sequences[session_id]
+        _session_sequences[session_id] = next_num + 1
+    return next_num