snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/utils/udf_helper.py

@@ -0,0 +1,328 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import base64
+import inspect
+import json
+import sys
+from typing import NamedTuple, Optional
+
+import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
+from pyspark.errors.exceptions.base import AnalysisException
+
+import snowflake.snowpark.functions as snowpark_fn
+import snowflake.snowpark_connect.tcm as tcm
+import snowflake.snowpark_connect.utils.udf_utils as udf_utils
+from snowflake.snowpark import Column, Session
+from snowflake.snowpark.types import DataType, _parse_datatype_json_value
+from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
+from snowflake.snowpark_connect.config import global_config
+from snowflake.snowpark_connect.expression.map_expression import (
+    map_single_column_expression,
+)
+from snowflake.snowpark_connect.expression.typer import ExpressionTyper
+from snowflake.snowpark_connect.utils.context import (
+    get_is_aggregate_function,
+    get_is_evaluating_join_condition,
+)
+from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
+from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
+
+CREATE_UDF_SPROC_NAME_PREFIX = "__SC_BUILD_IN_CREATE_UDF"
+TEST_FLAG_FORCE_CREATE_SPROC = False
+
+
+class SnowparkUDF(NamedTuple):
+    name: str
+    return_type: DataType
+    input_types: list[DataType]
+    original_return_type: DataType | None
+
+
+def require_creating_udf_in_sproc(
+    udf_proto: expressions_proto.CommonInlineUserDefinedFunction,
+) -> bool:
+    """
+    If a python UDF's python version does not match with current runtime, create it in a sproc using that python version.
+    """
+    return (
+        TEST_FLAG_FORCE_CREATE_SPROC
+        or tcm.TCM_MODE
+        or (
+            udf_proto.WhichOneof("function") == "python_udf"
+            and udf_proto.python_udf.python_ver is not None
+            and f"{sys.version_info.major}.{sys.version_info.minor}"
+            != udf_proto.python_udf.python_ver
+        )
+    )
+
+
+def process_udf_in_sproc(
+    common_inline_user_defined_function: expressions_proto.CommonInlineUserDefinedFunction,
+    called_from: str,
+    return_type: DataType,
+    input_types: list | None = None,
+    input_column_names: list[str] | None = None,
+    udf_name: str | None = None,
+    replace: bool = False,
+    udf_packages: str = "",
+    udf_imports: str = "",
+    original_return_type: DataType | None = None,
+) -> SnowparkUDF:
+    """Helper method to call the sproc to create inline UDF and return the essential info of the UDF."""
+    session = get_or_create_snowpark_session()
+    sproc_name = _get_or_create_udf_sproc_helper(
+        session,
+        common_inline_user_defined_function.python_udf.python_ver,
+    )
+    udf_proto_encoded = base64.b64encode(
+        common_inline_user_defined_function.SerializeToString()
+    ).decode("ascii")
+
+    def gen_input_types_json_str(input_types: list | None = None) -> Optional[str]:
+        if input_types is None:
+            return None
+        return json.dumps([dt.json_value() for dt in input_types])
+
+    input_types_json_str = gen_input_types_json_str(input_types)
+
+    input_column_names_json_str = (
+        None if input_column_names is None else json.dumps(input_column_names)
+    )
+
+    original_return_type_json_str = (
+        None
+        if original_return_type is None
+        else json.dumps(original_return_type.json_value())
+    )
+
+    sproc_res = session.call(
+        sproc_name,
+        called_from,
+        json.dumps(return_type.json_value()),
+        input_types_json_str,
+        input_column_names_json_str,
+        snowpark_fn.lit(udf_name),
+        replace,
+        udf_packages,
+        udf_imports,
+        udf_proto_encoded,
+        original_return_type_json_str,
+    )
+
+    udf_attr = json.loads(sproc_res)
+    snowpark_udf = SnowparkUDF(
+        name=udf_attr["name"],
+        input_types=[_parse_datatype_json_value(t) for t in udf_attr["input_types"]],
+        return_type=_parse_datatype_json_value(udf_attr["return_type"]),
+        original_return_type=original_return_type,
+    )
+    if called_from == "register_udf":
+        session._udfs[
+            common_inline_user_defined_function.function_name.lower()
+        ] = snowpark_udf
+    return snowpark_udf
+
+
+def _get_or_create_udf_sproc_helper(
+    session: Session,
+    python_version: str = f"{sys.version_info.major}.{sys.version_info.minor}",
+) -> str:
+    """
+    This helper method will get or create a sproc in targeted python version to create the python UDF. The sproc's
+    return value is a json string containing the UDF's name, input types, and, return type.
+    """
+    sproc_name = f"{CREATE_UDF_SPROC_NAME_PREFIX}_{python_version.replace('.', '_')}"
+    if sproc_name in session._sprocs:
+        return sproc_name
+
+    inline_udf_utils_py_code = inspect.getsource(udf_utils)
+
+    create_udf_sproc_sql = f"""
+CREATE OR REPLACE TEMPORARY PROCEDURE {sproc_name}(
+    called_from VARCHAR,
+    return_type_json_str VARCHAR,
+    input_types_json_str VARCHAR,
+    input_column_names_json_str VARCHAR,
+    udf_name VARCHAR,
+    replace BOOLEAN,
+    udf_packages VARCHAR,
+    udf_imports VARCHAR,
+    base64_str VARCHAR,
+    original_return_type VARCHAR
+)
+RETURNS STRING
+LANGUAGE PYTHON
+RUNTIME_VERSION = '{python_version}'
+PACKAGES = ('pyspark>=3.5.0,<4', 'cloudpickle', 'snowflake-snowpark-python==1.32.0', 'grpcio>=1.48.1')
+HANDLER = 'create'
+EXECUTE AS CALLER
+AS $$
+import cloudpickle
+import base64
+from pyspark.sql.connect.proto.expressions_pb2 import CommonInlineUserDefinedFunction
+import json
+from snowflake.snowpark.types import *
+from typing import Optional
+from snowflake.snowpark.types import _parse_datatype_json_value
+
+{inline_udf_utils_py_code}
+
+def parse_input_types(input_types_json_str) -> Optional[list[DataType]]:
+    if input_types_json_str is None:
+        return None
+    input_types_json = json.loads(input_types_json_str)
+    return [_parse_datatype_json_value(t) for t in input_types_json]
+
+def parse_return_type(return_type_json_str) -> Optional[DataType]:
+    return_type_json = json.loads(return_type_json_str)
+    result = _parse_datatype_json_value(return_type_json)
+    if isinstance(result, (ArrayType, MapType, StructType)):
+        result = result._as_nested()
+    return result
+
+
+def create(session, called_from, return_type_json_str, input_types_json_str, input_column_names_json_str, udf_name, replace, udf_packages, udf_imports, b64_str, original_return_type):
+    import snowflake.snowpark.context as context
+    context._use_structured_type_semantics = True
+    context._is_snowpark_connect_compatible_mode = True
+
+    restored_bytes = base64.b64decode(b64_str.encode('ascii'))
+    udf_proto = CommonInlineUserDefinedFunction()
+    udf_proto.ParseFromString(restored_bytes)
+    udf_processor = ProcessCommonInlineUserDefinedFunction(
+        udf_proto,
+        input_types=parse_input_types(input_types_json_str),
+        return_type=parse_return_type(return_type_json_str),
+        called_from=called_from,
+        input_column_names=None if input_column_names_json_str is None else json.loads(input_column_names_json_str),
+        udf_name=udf_name,
+        replace=replace,
+        udf_packages=udf_packages,
+        udf_imports=udf_imports,
+        original_return_type=parse_return_type(original_return_type) if original_return_type else None,
+    )
+    udf = udf_processor.create_udf()
+    return json.dumps({{"name": udf.name, "return_type": udf._return_type.json_value(), "input_types": [t.json_value() for t in udf._input_types]}})
+$$;
+"""
+    session.sql(create_udf_sproc_sql).collect()
+    session._sprocs.add(sproc_name)
+    logger.info(f"Procedure {sproc_name} created")
+    return sproc_name
+
+
+def udf_check(
+    udf_proto: expressions_proto.CommonInlineUserDefinedFunction,
+) -> None:
+    _check_supported_udf(udf_proto)
+    _aggregate_function_check(udf_proto)
+
+
+def _check_supported_udf(
+    udf_proto: expressions_proto.CommonInlineUserDefinedFunction,
+) -> None:
+    match udf_proto.WhichOneof("function"):
+        case "python_udf":
+            pass
+        case "java_udf":
+            from snowflake.snowpark_connect.utils.session import (
+                get_or_create_snowpark_session,
+            )
+
+            session = get_or_create_snowpark_session()
+            if udf_proto.java_udf.class_name not in session._cached_java_udfs:
+                raise AnalysisException(
+                    f"Can not load class {udf_proto.java_udf.class_name}"
+                )
+            else:
+                raise ValueError(
+                    "Function type java_udf not supported for common inline user-defined function"
+                )
+        case _ as function_type:
+            raise ValueError(
+                f"Function type {function_type} not supported for common inline user-defined function"
+            )
+
+
+def _aggregate_function_check(
+    udf_proto: expressions_proto.CommonInlineUserDefinedFunction,
+):
+    name, is_aggregate_function = get_is_aggregate_function()
+    if not udf_proto.deterministic and name != "default" and is_aggregate_function:
+        raise AnalysisException(
+            f"[AGGREGATE_FUNCTION_WITH_NONDETERMINISTIC_EXPRESSION] Non-deterministic expression {name}({udf_proto.function_name}) should not appear in the arguments of an aggregate function."
+        )
+
+
+def _join_checks(snowpark_udf_arg_names: list[str]):
+    is_evaluating_join_condition = get_is_evaluating_join_condition()
+    is_left_evaluable, is_right_evaluable = False, False
+
+    for snowpark_udf_arg_name in snowpark_udf_arg_names:
+        # UDFs can only reference EITHER the left OR the right side of the join but not both.
+        # Example: Assume left has column a and right has column b.
+        # lambda a: str(a) is fine because it will only reference either the left dataframe or the right dataframe.
+        # lambda a, b: a == b is not fine because it will reference both of the dataframes.
+        is_left_evaluable = (
+            is_left_evaluable
+            or snowpark_udf_arg_name in is_evaluating_join_condition[2]
+        )
+        is_right_evaluable = (
+            is_right_evaluable
+            or snowpark_udf_arg_name in is_evaluating_join_condition[3]
+        )
+        # Check for implicit cartesian product only on inner joins. If crossjoin is disabled, raise an exception.
+        if (
+            is_evaluating_join_condition[0] == "INNER"
+            and not global_config.spark_sql_crossJoin_enabled
+            and is_left_evaluable
+            and is_right_evaluable
+        ):
+            raise AnalysisException(
+                f"Detected implicit cartesian product for {is_evaluating_join_condition[0]} join between logical plans. \n"
+                f"Join condition is missing or trivial. \n"
+                f"Either: use the CROSS JOIN syntax to allow cartesian products between those relations, or; "
+                f"enable implicit cartesian products by setting the configuration variable spark.sql.crossJoin.enabled=True."
+            )
+        if (
+            is_evaluating_join_condition[0] != "INNER"
+            and is_evaluating_join_condition[1]
+            and is_left_evaluable
+            and is_right_evaluable
+        ):
+            raise AnalysisException(
+                f"[UNSUPPORTED_FEATURE.PYTHON_UDF_IN_ON_CLAUSE] The feature is not supported: "
+                f"Python UDF in the ON clause of a {is_evaluating_join_condition[0]} JOIN. "
+                f"In case of an INNNER JOIN consider rewriting to a CROSS JOIN with a WHERE clause."
+            )
+
+
+def infer_snowpark_arguments(
+    udf_proto: expressions_proto.CommonInlineUserDefinedFunction,
+    column_mapping: ColumnNameMap,
+    typer: ExpressionTyper,
+) -> tuple[list[str], list[Column]]:
+    snowpark_udf_args: list[Column] = []
+    snowpark_udf_arg_names: list[str] = []
+    for arg_exp in udf_proto.arguments:
+        (
+            snowpark_udf_arg_name,
+            snowpark_udf_arg,
+        ) = map_single_column_expression(arg_exp, column_mapping, typer)
+        snowpark_udf_arg = snowpark_udf_arg.col
+        snowpark_udf_args.append(snowpark_udf_arg)
+        snowpark_udf_arg_names.append(snowpark_udf_arg_name)
+    _join_checks(snowpark_udf_arg_names)
+    return snowpark_udf_arg_names, snowpark_udf_args
+
+
+def gen_input_types(
+    snowpark_udf_args: list[Column],
+    typer: ExpressionTyper,
+):
+    input_types = []
+    for udf_arg in snowpark_udf_args:
+        input_types.extend(typer.type(udf_arg))
+    return input_types

snowflake/snowpark_connect/utils/udf_utils.py

@@ -0,0 +1,310 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+# Note: this class will be loaded in a stored procedure to create a Python UDF with different Python version.
+# So its dependencies are restricted to pandas, snowpark, and, pyspark
+import functools
+import inspect
+from typing import Iterator
+
+import pandas
+import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
+from pyspark.serializers import CloudPickleSerializer
+
+import snowflake.snowpark.context as context
+import snowflake.snowpark.functions as snowpark_fn
+from snowflake.snowpark._internal.udf_utils import extract_return_input_types
+from snowflake.snowpark._internal.utils import TempObjectType
+from snowflake.snowpark.types import (
+    DataType,
+    PandasDataFrameType,
+    StructType,
+    VariantType,
+)
+
+MAP_IN_ARROW_EVAL_TYPE = 207
+
+
+def create_null_safe_wrapper(func):
+    """
+    Create a wrapper function that handles sqlNullWrapper objects by converting them to None.
+    This is needed because when arguments are cast to VariantType(), Snowflake wraps NULL values
+    in sqlNullWrapper objects instead of passing None to UDFs.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        # Convert sqlNullWrapper objects to None before calling the actual UDF
+        processed_args = [
+            None if getattr(arg, "is_sql_null", False) else arg for arg in args
+        ]
+        processed_kwargs = {
+            k: (None if getattr(v, "is_sql_null", False) else v)
+            for k, v in kwargs.items()
+        }
+        return func(*processed_args, **processed_kwargs)
+
+    return wrapper
+
+
+class ProcessCommonInlineUserDefinedFunction:
+    def __init__(
+        self,
+        common_inline_user_defined_function: expressions_proto.CommonInlineUserDefinedFunction,
+        called_from: str,
+        return_type: DataType,
+        input_column_names: list[str] | None = None,
+        input_types: list | None = None,
+        udf_name: str | None = None,
+        replace: bool = False,
+        udf_packages: str = "",
+        udf_imports: str = "",
+        original_return_type: DataType | None = None,
+    ) -> None:
+        context._use_structured_type_semantics = True
+        context._is_snowpark_connect_compatible_mode = True
+
+        self._function_name = common_inline_user_defined_function.function_name
+        self._is_deterministic = common_inline_user_defined_function.deterministic
+        self._arguments = common_inline_user_defined_function.arguments
+        self._function_type = common_inline_user_defined_function.WhichOneof("function")
+        self._input_types = input_types
+        self._udf_name = udf_name
+        self._replace = replace
+        self._called_from = called_from
+        self._input_column_names = input_column_names
+        self._udf_packages = udf_packages
+        self._udf_imports = udf_imports
+        self._original_return_type = original_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
+                )
+                self._command = common_inline_user_defined_function.python_udf.command
+                self._python_ver = (
+                    common_inline_user_defined_function.python_udf.python_ver
+                )
+            case _:
+                raise ValueError(
+                    f"Function type {self._function_type} not supported for common inline user-defined function"
+                )
+
+    @property
+    def snowpark_udf_args(self):
+        if self._column_mapping is not None:
+            return self._snowpark_udf_args
+        else:
+            raise ValueError(
+                "Column mapping is not provided, cannot get snowpark udf args"
+            )
+
+    @property
+    def snowpark_udf_arg_names(self):
+        if self._column_mapping is not None:
+            return self._snowpark_udf_arg_names
+        else:
+            raise ValueError(
+                "Column mapping is not provided, cannot get snowpark udf arg names"
+            )
+
+    def _create_python_udf(self):
+        def update_none_input_types():
+            if self._input_types is None:
+                # If any of the parameters don't have type hint, we use Snowflake Variant type.
+                func_signature = inspect.signature(callable_func)
+                self._input_types = [VariantType()] * len(func_signature.parameters)
+
+        (
+            callable_func,
+            pyspark_output_type,
+        ) = CloudPickleSerializer().loads(self._command)
+
+        # Wrap callable with a function which changes the current working directory
+        original_callable = callable_func
+
+        def import_staged_files():
+            import os
+            import shutil
+            import sys
+            import tarfile
+            import zipfile
+
+            # Change directory to the one containing the UDF imported files
+            import_path = sys._xoptions["snowflake_import_directory"]
+            tmp_path = f"/tmp/sas-{os.getpid()}"
+            os.makedirs(tmp_path, exist_ok=True)
+            os.chdir(tmp_path)
+            shutil.copytree(import_path, tmp_path, dirs_exist_ok=True)
+
+            # Extract all archives
+            # This has to be done inside the UDF because Snowflake prevents from loading multiple files with the same name
+            # even though they are in different paths.
+            archives = os.listdir(".")
+            for archive in archives:
+                if not archive.endswith(".archive"):
+                    # Skip files that are not archives
+                    continue
+                elif archive.endswith(".zip.archive") or archive.endswith(
+                    ".jar.archive"
+                ):
+                    with zipfile.ZipFile(archive, "r") as zip_ref:
+                        zip_ref.extractall(archive[: -len(".archive")])
+                elif archive.endswith(".tar.gz.archive") or archive.endswith(
+                    ".tgz.archive"
+                ):
+                    with tarfile.open(archive, "r:gz") as tar_ref:
+                        tar_ref.extractall(archive[: -len(".archive")])
+                elif archive.endswith(".tar.archive"):
+                    with tarfile.open(archive, "r") as tar_ref:
+                        tar_ref.extractall(archive[: -len(".archive")])
+                os.remove(archive)
+
+        def callable_func(*args, **kwargs):
+            import_staged_files()
+            return original_callable(*args, **kwargs)
+
+        callable_func.__signature__ = inspect.signature(original_callable)
+        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:
+            imports = []
+
+        update_none_input_types()
+
+        needs_struct_conversion = isinstance(self._original_return_type, StructType)
+
+        match self._called_from:
+            case "map_map_partitions":
+                column_names = self._input_column_names
+                eval_type = self._eval_type
+
+                def wrapped_function(*args):
+                    result = callable_func(
+                        pandas.DataFrame(iter([list(args)]), columns=column_names)
+                    )
+                    if eval_type == MAP_IN_ARROW_EVAL_TYPE:
+                        if not isinstance(result, Iterator) and not hasattr(
+                            result, "__iter__"
+                        ):
+                            raise RuntimeError(
+                                f"snowpark_connect::UDF_RETURN_TYPE Return type of the user-defined function should be "
+                                f"iterator of pyarrow.RecordBatch, but is {type(result).__name__}"
+                            )
+                        import pyarrow as pa
+
+                        for elem in result:
+                            if not isinstance(elem, pa.RecordBatch):
+                                raise RuntimeError(
+                                    f"snowpark_connect::UDF_RETURN_TYPE Return type of the user-defined function should "
+                                    f"be iterator of pyarrow.RecordBatch, but is iterator of {type(elem).__name__}"
+                                )
+                    return result
+
+                udf_function = create_null_safe_wrapper(wrapped_function)
+                packages += ["pyarrow", "pandas"]
+            case _:
+                if not needs_struct_conversion:
+                    return snowpark_fn.udf(
+                        create_null_safe_wrapper(callable_func),
+                        return_type=self._return_type,
+                        input_types=self._input_types,
+                        name=self._udf_name,
+                        replace=self._replace,
+                        packages=packages,
+                        imports=imports,
+                        immutable=self._is_deterministic,
+                    )
+
+                is_pandas_udf, _, return_types, _ = extract_return_input_types(
+                    callable_func,
+                    self._original_return_type,
+                    self._input_types,
+                    TempObjectType.FUNCTION,
+                )
+                if is_pandas_udf and isinstance(return_types, PandasDataFrameType):
+                    # Snowpark Python UDFs only support returning a Pandas Series.
+                    # We change the return type to make the input callable compatible with Snowpark Python UDFs,
+                    # and then in the wrapper function we convert the pandas DataFrame of the
+                    # original callable to a Pandas Series.
+                    original_callable.__annotations__["return"] = pandas.Series
+
+                field_names = [
+                    field.name for field in self._original_return_type.fields
+                ]
+
+                def struct_wrapper(*args):
+                    result = callable_func(*args)
+                    if isinstance(result, (tuple, list)):
+                        # Convert tuple/list to dict using struct field names
+                        if len(result) == len(field_names):
+                            return dict(zip(field_names, result))
+                    return result
+
+                def pandas_struct_wrapper(*args):
+                    # inspired by the following snowpark modin code to handle Pandas int/bool/null data in Snowflake VariantType
+                    # https://github.com/snowflakedb/snowpark-python/blob/e095d5a54f3a697416c3f1df87d239def47a5495/src/snowflake/snowpark/modin/plugin/_internal/apply_utils.py#L1309-L1366
+                    def convert_to_snowflake_compatible_type(value):
+                        import numpy as np
+                        from pandas.api.types import is_scalar
+
+                        if is_scalar(value) and pandas.isna(value):
+                            return None
+
+                        return (
+                            int(value)
+                            if np.issubdtype(type(value), np.integer)
+                            else (
+                                bool(value)
+                                if np.issubdtype(type(value), np.bool_)
+                                else value
+                            )
+                        )
+
+                    result = callable_func(*args)
+                    assert (
+                        len(result) == 1
+                    ), "Expected result to be a single row DataFrame"
+                    # df.applymap doesn't help here, the original type was preserved, hence we convert each value
+                    row_data = [
+                        convert_to_snowflake_compatible_type(value)
+                        for value in result.iloc[0].tolist()
+                    ]
+                    result = pandas.Series([dict(zip(field_names, row_data))])
+                    return result
+
+                if is_pandas_udf:
+                    udf_function = pandas_struct_wrapper
+                    if isinstance(return_types, PandasDataFrameType):
+                        udf_function.__annotations__ = original_callable.__annotations__
+                else:
+                    udf_function = create_null_safe_wrapper(struct_wrapper)
+
+        return snowpark_fn.udf(
+            udf_function,
+            return_type=self._return_type,
+            input_types=self._input_types,
+            name=self._udf_name,
+            replace=self._replace,
+            packages=packages,
+            imports=imports,
+            immutable=self._is_deterministic,
+        )
+
+    def create_udf(self):
+        match self._function_type:
+            case "python_udf":
+                return self._create_python_udf()
+            case _:
+                raise ValueError(
+                    f"Function type {self._function_type} not supported for common inline user-defined function"
+                )