SNOW-3152139: Create UDF_init_once decorator in python client (#4269)

Author: sfc-gh-nirsharma

CHANGELOG.md

@@ -4,6 +4,10 @@
 
 ### Snowpark Python API Updates
 
+#### New Features
+
+- Added the `udf_init_once` decorator in `snowflake.snowpark.functions` for marking functions to be executed once during pre-fork initialization on Snowflake workers, matching the server-side `_snowflake.udf_init_once` API.
+
 #### Bug Fixes
 
 - Fixed a bug where stage paths and file format names that contain single quotes were not consistently escaped when generating SQL, which could produce malformed statements. This affects `INFER_SCHEMA` (used by `DataFrameReader.csv`/`json`/`parquet`/`orc`/`avro`) and `COPY FILES` (used by `FileOperation.copy_files`).

docs/source/snowpark/functions.rst

@@ -575,6 +575,7 @@ Functions
     typeof
     udaf
     udf
+    udf_init_once
     udtf
     unbase64
     unicode

src/snowflake/snowpark/functions.py

@@ -158,6 +158,7 @@
     <BLANKLINE>
 """
 import functools
+import inspect
 import typing
 from functools import reduce
 import json  # noqa: F401  (only referenced by doctests in this module)
@@ -10329,6 +10330,64 @@ def _inner(*args, **kwargs):
     return _decorator
 
 
+def udf_init_once(func: Callable) -> Callable:
+    """Mark a function to be executed once during pre-fork initialization.
+
+    This decorator has the same signature as the server-side ``_snowflake.udf_init_once``.
+    It is a no-op for local invocation. Functions decorated with ``@udf_init_once`` run in
+    the head worker process before individual workers are forked. The initialized state
+    (e.g., loaded models, computed lookup tables) is shared across all workers via
+    Copy-On-Write.
+
+    The decorated function must:
+        - Accept zero arguments
+        - Be a callable (function, lambda, or callable object)
+
+    Multiple ``@udf_init_once`` functions are executed in the order they are defined.
+
+    Example::
+
+        from snowflake.snowpark.functions import udf_init_once
+
+        model = None
+
+        @udf_init_once
+        def load_model():
+            global model
+            model = 42
+
+    Use this decorator in handler files registered via
+    :meth:`~snowflake.snowpark.udf.UDFRegistration.register_from_file`.
+    On the Snowflake server the ``_snowflake.udf_init_once`` implementation
+    is used instead; this client-side definition provides the same API so
+    that handler files can be tested locally.
+
+    Args:
+        func: The init function to decorate. Must be callable and accept zero arguments.
+
+    Returns:
+        The decorated function, unchanged.
+
+    See Also:
+        - :func:`udf`
+        - :meth:`~snowflake.snowpark.udf.UDFRegistration.register_from_file`
+    """
+    if not callable(func):
+        raise TypeError(
+            f"@udf_init_once target must be callable, got {type(func).__name__}"
+        )
+    sig = inspect.signature(func)
+    if len(sig.parameters) != 0:
+        raise TypeError(
+            f"@udf_init_once function must take 0 arguments, got {len(sig.parameters)}"
+        )
+    # This client-side decorator is a no-op that returns the function unchanged;
+    # it exists only to mirror the ``_snowflake.udf_init_once`` signature so that
+    # handler files import and validate cleanly when authored/tested locally. The
+    # actual pre-fork execution is performed server-side by ``_snowflake``.
+    return func
+
+
 @publicapi
 def pandas_udtf(
     handler: Optional[Callable] = None,

tests/integ/test_function.py

@@ -167,6 +167,7 @@
 import functools
 from snowflake.connector.options import installed_pandas
 from snowflake.snowpark.functions import udf, vectorized
+from snowflake.snowpark.udf import UserDefinedFunction
 from snowflake.snowpark.types import (
     ArrayType,
     BooleanType,
@@ -2710,3 +2711,28 @@ def add_series(s1, s2):
     )
     res = df.select(add_series("a", "b").alias("result")).collect()
     assert [row.RESULT for row in res] == [3, 30, 300, 15]
+
+
+@pytest.mark.skipif(
+    "config.getoption('local_testing_mode', default=False)",
+    reason="UDF init_once is not supported in Local Testing",
+)
+def test_udf_init_once_register_from_file(session):
+    """A handler file using @udf_init_once registers successfully via register_from_file.
+
+    ``@udf_init_once`` is a client-side API-parity shim: it validates the target
+    and returns it unchanged, mirroring the ``_snowflake.udf_init_once`` signature
+    so handler files import cleanly both locally and on the server. Server-side
+    pre-fork execution of the decorated init function is not yet GA, so this test
+    only verifies that a handler file importing ``udf_init_once`` can be registered
+    via ``register_from_file``; it intentionally does not invoke the UDF to assert
+    the init function's effect.
+    """
+    multiply_udf = session.udf.register_from_file(
+        file_path="tests/resources/test_udf_dir/test_udf_init_once_file.py",
+        func_name="multiply",
+        return_type=IntegerType(),
+        input_types=[IntegerType()],
+    )
+    assert isinstance(multiply_udf, UserDefinedFunction)
+    assert multiply_udf.name

tests/resources/test_udf_dir/test_udf_init_once_file.py

@@ -0,0 +1,17 @@
+try:
+    from _snowflake import udf_init_once
+except ModuleNotFoundError:
+    from snowflake.snowpark.functions import udf_init_once
+
+
+_multiplier = 1
+
+
+@udf_init_once
+def setup():
+    global _multiplier
+    _multiplier = 10
+
+
+def multiply(x):
+    return x * _multiplier

tests/unit/scala/test_utils_suite.py

@@ -19,6 +19,7 @@
     get_stage_parts,
     get_temp_type_for_object,
     get_udf_upload_prefix,
+    is_cloud_path,
     is_snowflake_quoted_id_case_insensitive,
     is_snowflake_unquoted_suffix_case_insensitive,
     is_sql_select_statement,
@@ -30,7 +31,6 @@
     validate_object_name,
     warning,
     zip_file_or_directory_to_stream,
-    is_cloud_path,
 )
 from tests.utils import IS_WINDOWS, TestFiles
 
@@ -105,32 +105,32 @@ def test_calculate_checksum():
     if IS_WINDOWS:
         assert (
             calculate_checksum(test_files.test_udf_directory)
-            == "7ff46b5f3765187c7355852811c92ff232feb2f7207a12f1ab6d0b921319643b"
+            == "a72518e02fe61d2bfb6e54aa684b9af2da7845f4fb1a758c3c3974d35bf58d8f"
         )
         assert (
             calculate_checksum(test_files.test_udf_directory, algorithm="md5")
-            == "f6c1984af9ece1bd68edf16ae1a7f992"
+            == "34fa7ef3aa8ff5c19b328fcc8f601e1e"
         )
     else:
         assert (
             calculate_checksum(test_files.test_udf_directory)
-            == "3a2607ef293801f59e7840f5be423d4a55edfe2ac732775dcfda01205df377f0"
+            == "d472060e6d717517a3f9c7048ac43fc0c646467a6b3772f63355071a65ad8ecf"
         )
         assert (
             calculate_checksum(test_files.test_udf_directory, algorithm="md5")
-            == "b72b61c8d5639fff8aa9a80278dba60f"
+            == "322ad29c4018a1375f61b670196cf902"
         )
         # Validate that hashes are different when reading whole dir.
         # Using a sufficiently small chunk size so that the hashes differ.
         assert (
             calculate_checksum(test_files.test_udf_directory, chunk_size=128)
-            == "c071de824a67c083edad45c2b18729e17c50f1b13be980140437063842ea2469"
+            == "40d4811752a978779518082f0312f8823cb46dd84d12c68a7bb456c388f38df4"
         )
         assert (
             calculate_checksum(
                 test_files.test_udf_directory, chunk_size=128, whole_file_hash=True
             )
-            == "3a2607ef293801f59e7840f5be423d4a55edfe2ac732775dcfda01205df377f0"
+            == "d472060e6d717517a3f9c7048ac43fc0c646467a6b3772f63355071a65ad8ecf"
         )
 
 
@@ -258,6 +258,7 @@ def check_zip_files_and_close_stream(input_stream, expected_files):
             [
                 "test_udf_dir/",
                 "test_udf_dir/test_another_udf_file.py",
+                "test_udf_dir/test_udf_init_once_file.py",
                 "test_udf_dir/test_pandas_udf_file.py",
                 "test_udf_dir/test_udf_file.py",
             ],
@@ -272,6 +273,7 @@ def check_zip_files_and_close_stream(input_stream, expected_files):
             [
                 "test_udf_dir/",
                 "test_udf_dir/test_another_udf_file.py",
+                "test_udf_dir/test_udf_init_once_file.py",
                 "test_udf_dir/test_pandas_udf_file.py",
                 "test_udf_dir/test_udf_file.py",
             ],
@@ -287,6 +289,7 @@ def check_zip_files_and_close_stream(input_stream, expected_files):
                 "resources/",
                 "resources/test_udf_dir/",
                 "resources/test_udf_dir/test_another_udf_file.py",
+                "resources/test_udf_dir/test_udf_init_once_file.py",
                 "resources/test_udf_dir/test_pandas_udf_file.py",
                 "resources/test_udf_dir/test_udf_file.py",
             ],
@@ -386,6 +389,7 @@ def check_zip_files_and_close_stream(input_stream, expected_files):
                 "resources/test_sp_dir/test_table_sp_file.py",
                 "resources/test_udf_dir/",
                 "resources/test_udf_dir/test_another_udf_file.py",
+                "resources/test_udf_dir/test_udf_init_once_file.py",
                 "resources/test_udf_dir/test_pandas_udf_file.py",
                 "resources/test_udf_dir/test_udf_file.py",
                 "resources/test_udtf_dir/",

tests/unit/test_udf_utils.py

@@ -82,6 +82,36 @@ def test_generate_python_code_exception():
         assert "Source code comment could not be generated" in generated_code
 
 
+def test_udf_init_once_decorator_simple():
+    """Test @udf_init_once as simple decorator (same API as _snowflake.udf_init_once)."""
+    from snowflake.snowpark.functions import udf_init_once
+
+    @udf_init_once
+    def my_init():
+        pass
+
+    # The client-side decorator is a no-op: it returns the function unchanged.
+    assert my_init.__name__ == "my_init"
+    assert callable(my_init)
+
+
+def test_udf_init_once_rejects_non_zero_arg_function():
+    from snowflake.snowpark.functions import udf_init_once
+
+    with pytest.raises(TypeError, match="must take 0 arguments"):
+
+        @udf_init_once
+        def bad_init(x):
+            return x
+
+
+def test_udf_init_once_rejects_non_callable():
+    from snowflake.snowpark.functions import udf_init_once
+
+    with pytest.raises(TypeError, match="must be callable"):
+        udf_init_once("not_a_function")
+
+
 @pytest.mark.parametrize(
     "packages",
     [