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

snowflake/snowpark_connect/utils/java_udtf_utils.py

@@ -0,0 +1,239 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import hashlib
+from dataclasses import dataclass
+
+from pyspark.sql.connect.proto.expressions_pb2 import CommonInlineUserDefinedFunction
+
+from snowflake.snowpark.types import ArrayType, MapType, VariantType
+from snowflake.snowpark_connect.resources_initializer import (
+    ensure_scala_udf_jars_uploaded,
+)
+from snowflake.snowpark_connect.type_mapping import (
+    map_type_to_snowflake_type,
+    proto_to_snowpark_type,
+)
+from snowflake.snowpark_connect.utils.jvm_udf_utils import (
+    NullHandling,
+    Param,
+    ReturnType,
+    Signature,
+    build_jvm_udxf_imports,
+    map_type_to_java_type,
+)
+from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
+
+JAVA_UDTF_PREFIX = "__SC_JAVA_UDTF_"
+
+SCALA_INPUT_VARIANT = """
+Object mappedInput = com.snowflake.sas.scala.UdfPacketUtils$.MODULE$.fromVariant(udfPacket, input, 0);
+
+java.util.Iterator<Object> javaInput = Arrays.asList(mappedInput).iterator();
+scala.collection.Iterator<Object> scalaInput = new scala.collection.AbstractIterator<Object>() {
+    public boolean hasNext() { return javaInput.hasNext(); }
+    public Object next() { return javaInput.next(); }
+};
+"""
+
+SCALA_INPUT_SIMPLE_TYPE = """
+java.util.Iterator<__iterator_type__> javaInput = Arrays.asList(input).iterator();
+scala.collection.Iterator<__iterator_type__> scalaInput = new scala.collection.AbstractIterator<__iterator_type__>() {
+    public boolean hasNext() { return javaInput.hasNext(); }
+    public __iterator_type__ next() { return javaInput.next(); }
+};
+"""
+
+UDTF_TEMPLATE = """
+import org.apache.spark.sql.connect.common.UdfPacket;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.ObjectInputStream;
+import java.io.Serializable;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+
+import java.util.*;
+import java.lang.*;
+import java.util.stream.Collectors;
+import com.snowflake.snowpark_java.types.*;
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
+
+public class OutputRow {
+  public Variant __java_udtf_prefix__C1;
+  public OutputRow(Variant __java_udtf_prefix__C1) {
+    this.__java_udtf_prefix__C1 = __java_udtf_prefix__C1;
+  }
+}
+
+public class JavaUdtfHandler {
+    private final static String OPERATION_FILE = "__operation_file__";
+    private static scala.Function1<scala.collection.Iterator<__iterator_type__>, scala.collection.Iterator<Object>> operation = null;
+    private static UdfPacket udfPacket = null;
+
+  public static Class getOutputClass() { return OutputRow.class; }
+
+    private static void loadOperation() throws IOException, ClassNotFoundException {
+        if (operation != null) {
+            return; // Already loaded
+        }
+
+        udfPacket = com.snowflake.sas.scala.Utils$.MODULE$.deserializeUdfPacket(OPERATION_FILE);
+        operation = (scala.Function1<scala.collection.Iterator<__iterator_type__>, scala.collection.Iterator<Object>>) udfPacket.function();
+    }
+
+  public Stream<OutputRow> process(__input_type__ input) throws IOException, ClassNotFoundException {
+        loadOperation();
+
+        __scala_input__
+
+        scala.collection.Iterator<Object> scalaResult = operation.apply(scalaInput);
+
+        java.util.Iterator<Variant> javaResult = new java.util.Iterator<Variant>() {
+            public boolean hasNext() { return scalaResult.hasNext(); }
+            public Variant next() {
+                return com.snowflake.sas.scala.Utils$.MODULE$.toVariant(scalaResult.next(), udfPacket);
+            }
+        };
+
+        return StreamSupport.stream(Spliterators.spliteratorUnknownSize(javaResult, Spliterator.ORDERED), false)
+                .map(i -> new OutputRow(i));
+  }
+
+  public Stream<OutputRow> endPartition() {
+    return Stream.empty();
+  }
+}
+"""
+
+
+@dataclass(frozen=True)
+class JavaUDTFDef:
+    """
+    Complete definition for creating a Java UDTF in Snowflake.
+
+    Contains all the information needed to generate the CREATE FUNCTION SQL statement
+    and the Java code body for the UDTF.
+
+    Attributes:
+        name: UDTF name
+        signature: SQL signature (for Snowflake function definition)
+        java_signature: Java signature (for Java code generation)
+        imports: List of JAR files to import
+        null_handling: Null handling behavior (defaults to RETURNS_NULL_ON_NULL_INPUT)
+    """
+
+    name: str
+    signature: Signature
+    java_signature: Signature
+    imports: list[str]
+    null_handling: NullHandling = NullHandling.RETURNS_NULL_ON_NULL_INPUT
+
+    def _gen_body_java(self) -> str:
+        returns_variant = self.signature.returns.data_type == "VARIANT"
+        return_type = (
+            "Variant" if returns_variant else self.java_signature.returns.data_type
+        )
+
+        is_variant_input = self.java_signature.params[0].data_type.lower() == "variant"
+
+        scala_input_template = (
+            SCALA_INPUT_VARIANT if is_variant_input else SCALA_INPUT_SIMPLE_TYPE
+        )
+
+        iterator_type = (
+            "Object" if is_variant_input else self.java_signature.params[0].data_type
+        )
+
+        return (
+            UDTF_TEMPLATE.replace("__operation_file__", self.imports[0].split("/")[-1])
+            .replace("__scala_input__", scala_input_template)
+            .replace("__iterator_type__", iterator_type)
+            .replace("__input_type__", self.java_signature.params[0].data_type)
+            .replace("__return_type__", return_type)
+            .replace("__java_udtf_prefix__", JAVA_UDTF_PREFIX)
+        )
+
+    def to_create_function_sql(self) -> str:
+        args = ", ".join(
+            [f"{param.name} {param.data_type}" for param in self.signature.params]
+        )
+
+        def quote_single(s: str) -> str:
+            """Helper function to wrap strings in single quotes for SQL."""
+            return "'" + s + "'"
+
+        # Handler and imports
+        imports_sql = f"IMPORTS = ({', '.join(quote_single(x) for x in self.imports)})"
+
+        return f"""
+create or replace function {self.name}({args})
+returns table ({JAVA_UDTF_PREFIX}C1 VARIANT)
+language java
+runtime_version = 17
+PACKAGES = ('com.snowflake:snowpark:latest')
+{imports_sql}
+handler='JavaUdtfHandler'
+as
+$$
+{self._gen_body_java()}
+$$;"""
+
+
+def create_java_udtf_for_scala_flatmap_handling(
+    udf_proto: CommonInlineUserDefinedFunction,
+) -> str:
+    ensure_scala_udf_jars_uploaded()
+
+    return_type = proto_to_snowpark_type(udf_proto.scalar_scala_udf.outputType)
+
+    session = get_or_create_snowpark_session()
+
+    return_type_java = map_type_to_java_type(return_type)
+    sql_return_type = map_type_to_snowflake_type(return_type)
+
+    java_input_params: list[Param] = []
+    sql_input_params: list[Param] = []
+    for i, input_type_proto in enumerate(udf_proto.scalar_scala_udf.inputTypes):
+        input_type = proto_to_snowpark_type(input_type_proto)
+
+        param_name = "arg" + str(i)
+
+        if isinstance(input_type, (ArrayType, MapType, VariantType)):
+            java_type = "Variant"
+            snowflake_type = "Variant"
+        else:
+            java_type = map_type_to_java_type(input_type)
+            snowflake_type = map_type_to_snowflake_type(input_type)
+
+        java_input_params.append(Param(param_name, java_type))
+        sql_input_params.append(Param(param_name, snowflake_type))
+
+    udtf_name = (
+        JAVA_UDTF_PREFIX + hashlib.md5(udf_proto.scalar_scala_udf.payload).hexdigest()
+    )
+
+    imports = build_jvm_udxf_imports(
+        session,
+        udf_proto.scalar_scala_udf.payload,
+        udtf_name,
+    )
+
+    udtf = JavaUDTFDef(
+        name=udtf_name,
+        signature=Signature(
+            params=sql_input_params, returns=ReturnType(sql_return_type)
+        ),
+        imports=imports,
+        java_signature=Signature(
+            params=java_input_params, returns=ReturnType(return_type_java)
+        ),
+    )
+
+    sql = udtf.to_create_function_sql()
+    session.sql(sql).collect()
+
+    return udtf_name

snowflake/snowpark_connect/utils/jvm_udf_utils.py

@@ -0,0 +1,281 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+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 import snowpark
+from snowflake.snowpark_connect.config import get_scala_version
+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.resources_initializer import (
+    JSON_4S_JAR_212,
+    JSON_4S_JAR_213,
+    RESOURCE_PATH,
+    SAS_SCALA_UDF_JAR_212,
+    SAS_SCALA_UDF_JAR_213,
+    SCALA_REFLECT_JAR_212,
+    SCALA_REFLECT_JAR_213,
+    SPARK_COMMON_UTILS_JAR_212,
+    SPARK_COMMON_UTILS_JAR_213,
+    SPARK_CONNECT_CLIENT_JAR_212,
+    SPARK_CONNECT_CLIENT_JAR_213,
+    SPARK_SQL_JAR_212,
+    SPARK_SQL_JAR_213,
+)
+
+
+@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
+
+
+def build_jvm_udxf_imports(
+    session: snowpark.Session, payload: bytes, udf_name: str
+) -> List[str]:
+    """
+    Build the list of imports needed for the JVM UDxF.
+
+    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 UDxF
+
+    Args:
+        session: Snowpark session
+        payload: Binary payload containing the serialized Scala UDF
+        udf_name: Name of the Scala UDF (used for the binary file name)
+        is_map_return: Indicates if the UDxF returns a Map (affects imports)
+
+    Returns:
+        List of JAR file paths to be imported by the UDxF
+    """
+    # 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 + "/scala/bin/" + udf_name + ".bin"
+    session.file.put_stream(
+        payload_as_stream,
+        closure_binary_file,
+        overwrite=True,
+    )
+
+    # Format the user jars to be used in the IMPORTS clause of the stored procedure.
+    return (
+        [closure_binary_file]
+        + _scala_static_imports_for_udf(stage_resource_path)
+        + list(session._artifact_jars)
+    )
+
+
+def _scala_static_imports_for_udf(stage_resource_path: str) -> list[str]:
+    scala_version = get_scala_version()
+    if scala_version == "2.12":
+        return [
+            f"{stage_resource_path}/{SPARK_CONNECT_CLIENT_JAR_212}",
+            f"{stage_resource_path}/{SPARK_COMMON_UTILS_JAR_212}",
+            f"{stage_resource_path}/{SPARK_SQL_JAR_212}",
+            f"{stage_resource_path}/{JSON_4S_JAR_212}",
+            f"{stage_resource_path}/{SAS_SCALA_UDF_JAR_212}",
+            f"{stage_resource_path}/{SCALA_REFLECT_JAR_212}",  # Required for deserializing Scala lambdas
+        ]
+
+    if scala_version == "2.13":
+        return [
+            f"{stage_resource_path}/{SPARK_CONNECT_CLIENT_JAR_213}",
+            f"{stage_resource_path}/{SPARK_COMMON_UTILS_JAR_213}",
+            f"{stage_resource_path}/{SPARK_SQL_JAR_213}",
+            f"{stage_resource_path}/{JSON_4S_JAR_213}",
+            f"{stage_resource_path}/{SAS_SCALA_UDF_JAR_213}",
+            f"{stage_resource_path}/{SCALA_REFLECT_JAR_213}",  # Required for deserializing Scala lambdas
+        ]
+
+    # invalid Scala version
+    exception = ValueError(
+        f"Unsupported Scala version: {scala_version}. Snowpark Connect supports Scala 2.12 and 2.13"
+    )
+    attach_custom_error_code(exception, ErrorCodes.INVALID_CONFIG_VALUE)
+    raise exception
+
+
+def map_type_to_java_type(
+    t: Union[snowpark_type.DataType, types_proto.DataType]
+) -> str:
+    """Maps a Snowpark or Spark protobuf type to a Java type string."""
+    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"{map_type_to_java_type(t.element_type)}[]"
+                if is_snowpark_type
+                else f"{map_type_to_java_type(t.array.element_type)}[]"
+            )
+        case snowpark_type.BinaryType | "binary":
+            return "byte[]"
+        case snowpark_type.BooleanType | "boolean":
+            return "Boolean"
+        case snowpark_type.ByteType | "byte":
+            return "Byte"
+        case snowpark_type.DateType | "date":
+            return "java.sql.Date"
+        case snowpark_type.DecimalType | "decimal":
+            return "java.math.BigDecimal"
+        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 "Integer"
+        case snowpark_type.LongType | "long":
+            return "Long"
+        case snowpark_type.MapType | "map":  # can also map to OBJECT in Snowflake
+            key_type = (
+                map_type_to_java_type(t.key_type)
+                if is_snowpark_type
+                else map_type_to_java_type(t.map.key_type)
+            )
+            value_type = (
+                map_type_to_java_type(t.value_type)
+                if is_snowpark_type
+                else map_type_to_java_type(t.map.value_type)
+            )
+            return f"Map<{key_type}, {value_type}>"
+        case snowpark_type.NullType | "null":
+            return "String"  # cannot set the return type to Null in Snowpark Java UDAFs
+        case snowpark_type.ShortType | "short":
+            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 _:
+            exception = ValueError(f"Unsupported Snowpark type: {t}")
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+            raise exception
+
+
+def cast_java_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 or Struct, cast the argument arg_name to the correct type in Java."""
+    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 Java."""
+        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 f"Boolean.valueOf({arg_name}"
+            case snowpark_type.ByteType | "byte":
+                return arg_name + ".getBytes()[0]"  # 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 f"Double.valueOf({arg_name}"
+            case snowpark_type.FloatType | "float":
+                return f"Float.valueOf({arg_name}"
+            case snowpark_type.IntegerType | "integer":
+                return f"Integer.valueOf({arg_name}"
+            case snowpark_type.LongType | "long":
+                return f"Long.valueOf({arg_name}"
+            case snowpark_type.ShortType | "short":
+                return f"Short.valueOf({arg_name}"
+            case snowpark_type.StringType | "string" | "char" | "varchar":
+                return arg_name
+            case snowpark_type.TimestampType | "timestamp" | "timestamp_ntz":
+                return f"java.sql.Timestamp.valueOf({arg_name})"  # todo add test
+            case _:
+                exception = ValueError(f"Unsupported Snowpark type: {t}")
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+                raise exception
+
+    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
+        )
+        key_converter = "{" + convert_from_string_to_type("e.getKey()", key_type) + "}"
+        value_converter = (
+            "{" + convert_from_string_to_type("e.getValue()", value_type) + "}"
+        )
+        return f"""
+        {arg_name}.entrySet()
+            .stream()
+            .collect(Collectors.toMap(
+                e -> {key_converter},
+                e -> {value_converter}
+        ));
+        """
+    else:
+        return arg_name