snowpark-connect 0.24.0__py3-none-any.whl → 0.25.0__py3-none-any.whl

snowflake/snowpark_connect/utils/udf_utils.py

@@ -70,6 +70,14 @@ class ProcessCommonInlineUserDefinedFunction:
         self._arguments = common_inline_user_defined_function.arguments
         self._function_type = common_inline_user_defined_function.WhichOneof("function")
         self._input_types = input_types
+        if self._function_type == "scalar_scala_udf" and self._input_types is None:
+            # If the Scala UDF is created via `spark.udf.register`, the input types are not automatically inferred.
+            # Pass on the input types from the proto message to the Scala UDF handler.
+            self._scala_input_types = (
+                common_inline_user_defined_function.scalar_scala_udf.inputTypes
+            )
+        else:
+            self._scala_input_types = None
         self._udf_name = udf_name
         self._replace = replace
         self._called_from = called_from
@@ -77,9 +85,9 @@ class ProcessCommonInlineUserDefinedFunction:
         self._udf_packages = udf_packages
         self._udf_imports = udf_imports
         self._original_return_type = original_return_type
+        self._return_type = return_type
         match self._function_type:
             case "python_udf":
-                self._return_type = return_type
                 self._eval_type = (
                     common_inline_user_defined_function.python_udf.eval_type
                 )
@@ -87,6 +95,13 @@ class ProcessCommonInlineUserDefinedFunction:
                 self._python_ver = (
                     common_inline_user_defined_function.python_udf.python_ver
                 )
+            case "scalar_scala_udf":
+                self._payload = (
+                    common_inline_user_defined_function.scalar_scala_udf.payload
+                )
+                self._nullable = (
+                    common_inline_user_defined_function.scalar_scala_udf.nullable
+                )
             case _:
                 raise ValueError(
                     f"Function type {self._function_type} not supported for common inline user-defined function"
@@ -170,10 +185,6 @@ class ProcessCommonInlineUserDefinedFunction:
         if hasattr(original_callable, "__annotations__"):
             callable_func.__annotations__ = original_callable.__annotations__
 
-        if self._udf_packages:
-            packages = [p.strip() for p in self._udf_packages.strip("[]").split(",")]
-        else:
-            packages = []
         if self._udf_imports:
             imports = [i.strip() for i in self._udf_imports.split(",") if i.strip()]
         else:
@@ -183,6 +194,7 @@ class ProcessCommonInlineUserDefinedFunction:
 
         needs_struct_conversion = isinstance(self._original_return_type, StructType)
 
+        additional_packages = None
         match self._called_from:
             case "map_map_partitions":
                 column_names = self._input_column_names
@@ -211,7 +223,7 @@ class ProcessCommonInlineUserDefinedFunction:
                     return result
 
                 udf_function = create_null_safe_wrapper(wrapped_function)
-                packages += ["pyarrow", "pandas"]
+                additional_packages = ["pyarrow", "pandas"]
             case _:
                 if not needs_struct_conversion:
                     return snowpark_fn.udf(
@@ -220,7 +232,9 @@ class ProcessCommonInlineUserDefinedFunction:
                         input_types=self._input_types,
                         name=self._udf_name,
                         replace=self._replace,
-                        packages=packages,
+                        # Packages should already have been added to the session with
+                        # "snowpark.connect.udf.packages".
+                        packages=None,
                         imports=imports,
                         immutable=self._is_deterministic,
                     )
@@ -289,6 +303,17 @@ class ProcessCommonInlineUserDefinedFunction:
                 else:
                     udf_function = create_null_safe_wrapper(struct_wrapper)
 
+        packages = None
+        if additional_packages:
+            if not self._udf_packages:
+                packages = additional_packages
+            else:
+                existing_packages = [
+                    p.strip() for p in self._udf_packages.strip("[]").split(",")
+                ]
+                if not set(additional_packages).issubset(set(existing_packages)):
+                    packages = existing_packages + additional_packages
+
         return snowpark_fn.udf(
             udf_function,
             return_type=self._return_type,
@@ -304,6 +329,12 @@ class ProcessCommonInlineUserDefinedFunction:
         match self._function_type:
             case "python_udf":
                 return self._create_python_udf()
+            case "scalar_scala_udf":
+                from snowflake.snowpark_connect.utils.scala_udf_utils import (
+                    create_scala_udf,
+                )
+
+                return create_scala_udf(self)
             case _:
                 raise ValueError(
                     f"Function type {self._function_type} not supported for common inline user-defined function"

snowflake/snowpark_connect/utils/udtf_utils.py

@@ -52,21 +52,35 @@ def create_udtf(
         if hasattr(callable_func, "terminate"):
             callable_func.end_partition = callable_func.terminate
 
-    def process_packages(packages_str: str) -> list[str]:
+    def process_packages(packages_str: str) -> (list[str], bool):
         packages = []
         if packages_str and len(packages_str) > 0:
             packages = [p.strip() for p in packages_str.strip("[]").split(",")]
+
+        has_non_session_packages = False
+
         # Include pyarrow in packages when using Arrow-enabled UDTF
         if is_arrow_enabled and "pyarrow" not in packages:
             packages += ["pyarrow"]
+            has_non_session_packages = True
 
         if "pyspark" not in packages:
             # need this to support table argument in UDTF.
             packages += ["pyspark"]
+            has_non_session_packages = True
+
+        return (
+            packages,
+            has_non_session_packages,
+        )
 
-        return packages
+    packages: str | None
+    packages, has_non_session_packages = process_packages(packages)
+    if not has_non_session_packages:
+        # Skip package validation in the client as all packages were already added to the session
+        # via "snowpark.connect.udf.packages".
+        packages = None
 
-    packages = process_packages(packages)
     imports = [i.strip() for i in imports.split(",") if i.strip()]
     match called_from:
         case "register_udtf":

snowflake/snowpark_connect/version.py

@@ -3,4 +3,4 @@
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
 
-VERSION = (0, 24, 0)
+VERSION = (0, 25, 0)

{snowpark_connect-0.24.0.dist-info → snowpark_connect-0.25.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowpark-connect
-Version: 0.24.0
+Version: 0.25.0
 Summary: Snowpark Connect for Spark
 Author: Snowflake, Inc
 License: Apache License, Version 2.0