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

snowflake/snowpark_connect/expression/map_udf.py

@@ -7,20 +7,23 @@ import pyspark.sql.connect.proto.types_pb2 as types_proto
 
 import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
-from snowflake.snowpark.types import MapType, StructType, VariantType
+from snowflake.snowpark.types import ArrayType, MapType, StructType, VariantType
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
 from snowflake.snowpark_connect.config import global_config
+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.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.context import get_grouping_by_scala_udf_key
 from snowflake.snowpark_connect.utils.external_udxf_cache import (
     cache_external_udf,
     get_external_udf_from_cache,
 )
+from snowflake.snowpark_connect.utils.java_stored_procedure import create_java_udf
 from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
 from snowflake.snowpark_connect.utils.udf_helper import (
     SnowparkUDF,
-    gen_input_types,
     infer_snowpark_arguments,
     process_udf_in_sproc,
     require_creating_udf_in_sproc,
@@ -53,8 +56,14 @@ def cache_external_udf_wrapper(from_register_udf: bool):
                         session._udfs[udf_proto.function_name.lower()] = cached_udf
                     case "python_udf":
                         pass
+                    case "java_udf":
+                        session._udfs[udf_proto.function_name.lower()] = cached_udf
                     case _:
-                        raise ValueError(f"Unsupported UDF type: {function_type}")
+                        exception = ValueError(f"Unsupported UDF type: {function_type}")
+                        attach_custom_error_code(
+                            exception, ErrorCodes.UNSUPPORTED_OPERATION
+                        )
+                        raise exception
 
                 return cached_udf
 
@@ -94,13 +103,43 @@ def register_udf(
     match udf_proto.WhichOneof("function"):
         case "python_udf":
             output_type = udf_proto.python_udf.output_type
+            processed_return_type, original_return_type = process_udf_return_type(
+                output_type
+            )
         case "scalar_scala_udf":
+            # For Scala UDFs, always use VariantType as the processed type since all Scala UDFs
+            # return Variant. The actual type conversion happens after the UDF call.
             output_type = udf_proto.scalar_scala_udf.outputType
+            original_return_type = proto_to_snowpark_type(output_type)
+            processed_return_type = VariantType()
+        case "java_udf":
+            has_output_type = udf_proto.java_udf.HasField("output_type")
+            session = get_or_create_snowpark_session()
+            java_udf = create_java_udf(
+                session,
+                udf_proto.function_name,
+                udf_proto.java_udf.class_name,
+            )
+            original_return_type = java_udf._return_type
+            if has_output_type:
+                original_return_type = proto_to_snowpark_type(
+                    udf_proto.java_udf.output_type
+                )
+            udf = SnowparkUDF(
+                name=java_udf.name,
+                input_types=java_udf._input_types,
+                return_type=java_udf._return_type,
+                original_return_type=original_return_type,
+                cast_to_original_return_type=True,
+            )
+            session._udfs[udf_proto.function_name.lower()] = udf
+            return udf
         case _:
-            raise ValueError(
+            exception = ValueError(
                 f"Unsupported UDF type: {udf_proto.WhichOneof('function')}"
             )
-    processed_return_type, original_return_type = process_udf_return_type(output_type)
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
     session = get_or_create_snowpark_session()
     kwargs = {
         "common_inline_user_defined_function": udf_proto,
@@ -116,11 +155,15 @@ def register_udf(
     else:
         udf_processor = ProcessCommonInlineUserDefinedFunction(**kwargs)
         udf = udf_processor.create_udf()
+        is_scala_udf = udf_proto.WhichOneof("function") == "scalar_scala_udf"
+
         udf = SnowparkUDF(
             name=udf.name,
             input_types=udf._input_types,
             return_type=udf._return_type,
             original_return_type=original_return_type,
+            cast_to_original_return_type=is_scala_udf
+            or udf._return_type == VariantType(),
         )
         session._udfs[udf_proto.function_name.lower()] = udf
         # scala udfs can be also accessed using `udf.name`
@@ -136,19 +179,22 @@ def map_common_inline_user_defined_udf(
 ) -> tuple[str, TypedColumn]:
     udf_proto = exp.common_inline_user_defined_function
     udf_check(udf_proto)
-    snowpark_udf_arg_names, snowpark_udf_args = infer_snowpark_arguments(
+    snowpark_udf_arg_names, snowpark_udf_typed_args = infer_snowpark_arguments(
         udf_proto, column_mapping, typer
     )
-    input_types = gen_input_types(snowpark_udf_args, typer)
+    input_types = [a.typ for a in snowpark_udf_typed_args]
     match udf_proto.WhichOneof("function"):
         case "python_udf":
             processed_return_type, original_return_type = process_udf_return_type(
                 udf_proto.python_udf.output_type
             )
         case "scalar_scala_udf":
-            processed_return_type, original_return_type = process_udf_return_type(
+            # For Scala UDFs, always use VariantType as the processed type since all Scala UDFs
+            # return Variant. The actual type conversion happens after the UDF call.
+            original_return_type = proto_to_snowpark_type(
                 udf_proto.scalar_scala_udf.outputType
             )
+            processed_return_type = VariantType()
 
     @cache_external_udf_wrapper(from_register_udf=False)
     def get_snowpark_udf(
@@ -178,24 +224,44 @@ def map_common_inline_user_defined_udf(
         return snowpark_udf
 
     snowpark_udf = get_snowpark_udf(udf_proto)
-    udf_call_expr = snowpark_fn.call_udf(snowpark_udf.name, *snowpark_udf_args)
+    # Determine if we need to cast the result back to the original type
+    is_scala_udf = udf_proto.WhichOneof("function") == "scalar_scala_udf"
 
-    # If the original return type was MapType or StructType but we converted it to VariantType,
-    # we need to parse the JSON result back to the original type
-    if isinstance(original_return_type, (MapType, StructType)) and isinstance(
+    # For structured types (arrays, structs, maps), use to_variant instead of cast
+    # to ensure proper conversion to VARIANT type for Scala UDFS
+    converted_args = []
+    for tc in snowpark_udf_typed_args:
+        if is_scala_udf and isinstance(tc.typ, (ArrayType, StructType, MapType)):
+            converted_args.append(snowpark_fn.to_variant(tc.col))
+        else:
+            converted_args.append(tc.col)
+
+    udf_call_expr = snowpark_fn.call_udf(snowpark_udf.name, *converted_args)
+
+    # For Scala UDFs, always cast from Variant to the original type
+    # For Python UDFs, only cast if the original type was MapType or StructType
+    if is_scala_udf:
+        # All Scala UDFs return Variant, so we always need to cast back to the original type
+        result_expr = snowpark_fn.cast(udf_call_expr, original_return_type)
+        result_type = original_return_type
+
+    elif isinstance(original_return_type, (MapType, StructType)) and isinstance(
         processed_return_type, VariantType
     ):
-        # Parse JSON and cast back to original type
+        # Parse JSON and cast back to original type for Python UDFs
         result_expr = snowpark_fn.parse_json(udf_call_expr).cast(original_return_type)
         result_type = original_return_type
     else:
         result_expr = udf_call_expr
         result_type = snowpark_udf.return_type
 
-    return (
-        f"{udf_proto.function_name}({', '.join(snowpark_udf_arg_names)})",
-        TypedColumn(
-            result_expr,
-            lambda: [result_type],
-        ),
-    )
+    name = f"{udf_proto.function_name}({', '.join(snowpark_udf_arg_names)})"
+    if get_grouping_by_scala_udf_key() and not isinstance(
+        original_return_type, StructType
+    ):
+        name = (
+            "value"
+            if global_config.spark_sql_legacy_dataset_nameNonStructGroupingKeyAsValue
+            else "key"
+        )
+    return (name, TypedColumn(result_expr, lambda: [result_type]))