snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/includes/python/pyspark/sql/udtf.py

@@ -0,0 +1,325 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""
+User-defined table function related classes and functions
+"""
+import pickle
+import sys
+import warnings
+from typing import Any, Type, TYPE_CHECKING, Optional, Union
+
+from py4j.java_gateway import JavaObject
+
+from pyspark.errors import PySparkAttributeError, PySparkRuntimeError, PySparkTypeError
+from pyspark.rdd import PythonEvalType
+from pyspark.sql.column import _to_java_column, _to_seq
+from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version
+from pyspark.sql.types import StructType, _parse_datatype_string
+from pyspark.sql.udf import _wrap_function
+
+if TYPE_CHECKING:
+    from pyspark.sql._typing import ColumnOrName
+    from pyspark.sql.dataframe import DataFrame
+    from pyspark.sql.session import SparkSession
+
+__all__ = ["UDTFRegistration"]
+
+
+def _create_udtf(
+    cls: Type,
+    returnType: Union[StructType, str],
+    name: Optional[str] = None,
+    evalType: int = PythonEvalType.SQL_TABLE_UDF,
+    deterministic: bool = False,
+) -> "UserDefinedTableFunction":
+    """Create a Python UDTF with the given eval type."""
+    udtf_obj = UserDefinedTableFunction(
+        cls, returnType=returnType, name=name, evalType=evalType, deterministic=deterministic
+    )
+
+    return udtf_obj
+
+
+def _create_py_udtf(
+    cls: Type,
+    returnType: Union[StructType, str],
+    name: Optional[str] = None,
+    deterministic: bool = False,
+    useArrow: Optional[bool] = None,
+) -> "UserDefinedTableFunction":
+    """Create a regular or an Arrow-optimized Python UDTF."""
+    # Determine whether to create Arrow-optimized UDTFs.
+    if useArrow is not None:
+        arrow_enabled = useArrow
+    else:
+        from pyspark.sql import SparkSession
+
+        session = SparkSession._instantiatedSession
+        arrow_enabled = False
+        if session is not None:
+            value = session.conf.get("spark.sql.execution.pythonUDTF.arrow.enabled")
+            if isinstance(value, str) and value.lower() == "true":
+                arrow_enabled = True
+
+    eval_type: int = PythonEvalType.SQL_TABLE_UDF
+
+    if arrow_enabled:
+        # Return the regular UDTF if the required dependencies are not satisfied.
+        try:
+            require_minimum_pandas_version()
+            require_minimum_pyarrow_version()
+            eval_type = PythonEvalType.SQL_ARROW_TABLE_UDF
+        except ImportError as e:
+            warnings.warn(
+                f"Arrow optimization for Python UDTFs cannot be enabled: {str(e)}. "
+                f"Falling back to using regular Python UDTFs.",
+                UserWarning,
+            )
+
+    return _create_udtf(
+        cls=cls,
+        returnType=returnType,
+        name=name,
+        evalType=eval_type,
+        deterministic=deterministic,
+    )
+
+
+def _validate_udtf_handler(cls: Any) -> None:
+    """Validate the handler class of a UDTF."""
+
+    if not isinstance(cls, type):
+        raise PySparkTypeError(
+            error_class="INVALID_UDTF_HANDLER_TYPE", message_parameters={"type": type(cls).__name__}
+        )
+
+    if not hasattr(cls, "eval"):
+        raise PySparkAttributeError(
+            error_class="INVALID_UDTF_NO_EVAL", message_parameters={"name": cls.__name__}
+        )
+
+
+class UserDefinedTableFunction:
+    """
+    User-defined table function in Python
+
+    .. versionadded:: 3.5.0
+
+    Notes
+    -----
+    The constructor of this class is not supposed to be directly called.
+    Use :meth:`pyspark.sql.functions.udtf` to create this instance.
+
+    This API is evolving.
+    """
+
+    def __init__(
+        self,
+        func: Type,
+        returnType: Union[StructType, str],
+        name: Optional[str] = None,
+        evalType: int = PythonEvalType.SQL_TABLE_UDF,
+        deterministic: bool = False,
+    ):
+        _validate_udtf_handler(func)
+
+        self.func = func
+        self._returnType = returnType
+        self._returnType_placeholder: Optional[StructType] = None
+        self._inputTypes_placeholder = None
+        self._judtf_placeholder = None
+        self._name = name or func.__name__
+        self.evalType = evalType
+        self.deterministic = deterministic
+
+    @property
+    def returnType(self) -> StructType:
+        # `_parse_datatype_string` accesses to JVM for parsing a DDL formatted string.
+        # This makes sure this is called after SparkContext is initialized.
+        if self._returnType_placeholder is None:
+            if isinstance(self._returnType, str):
+                parsed = _parse_datatype_string(self._returnType)
+            else:
+                parsed = self._returnType
+            if not isinstance(parsed, StructType):
+                raise PySparkTypeError(
+                    error_class="UDTF_RETURN_TYPE_MISMATCH",
+                    message_parameters={
+                        "name": self._name,
+                        "return_type": f"{parsed}",
+                    },
+                )
+            self._returnType_placeholder = parsed
+        return self._returnType_placeholder
+
+    @property
+    def _judtf(self) -> JavaObject:
+        if self._judtf_placeholder is None:
+            self._judtf_placeholder = self._create_judtf(self.func)
+        return self._judtf_placeholder
+
+    def _create_judtf(self, func: Type) -> JavaObject:
+        from pyspark.sql import SparkSession
+
+        spark = SparkSession._getActiveSessionOrCreate()
+        sc = spark.sparkContext
+
+        try:
+            wrapped_func = _wrap_function(sc, func)
+        except pickle.PicklingError as e:
+            if "CONTEXT_ONLY_VALID_ON_DRIVER" in str(e):
+                raise PySparkRuntimeError(
+                    error_class="UDTF_SERIALIZATION_ERROR",
+                    message_parameters={
+                        "name": self._name,
+                        "message": "it appears that you are attempting to reference SparkSession "
+                        "inside a UDTF. SparkSession can only be used on the driver, "
+                        "not in code that runs on workers. Please remove the reference "
+                        "and try again.",
+                    },
+                ) from None
+            raise PySparkRuntimeError(
+                error_class="UDTF_SERIALIZATION_ERROR",
+                message_parameters={
+                    "name": self._name,
+                    "message": "Please check the stack trace and make sure the "
+                    "function is serializable.",
+                },
+            )
+
+        jdt = spark._jsparkSession.parseDataType(self.returnType.json())
+        assert sc._jvm is not None
+        judtf = sc._jvm.org.apache.spark.sql.execution.python.UserDefinedPythonTableFunction(
+            self._name, wrapped_func, jdt, self.evalType, self.deterministic
+        )
+        return judtf
+
+    def __call__(self, *cols: "ColumnOrName") -> "DataFrame":
+        from pyspark.sql import DataFrame, SparkSession
+
+        spark = SparkSession._getActiveSessionOrCreate()
+        sc = spark.sparkContext
+
+        judtf = self._judtf
+        jPythonUDTF = judtf.apply(spark._jsparkSession, _to_seq(sc, cols, _to_java_column))
+        return DataFrame(jPythonUDTF, spark)
+
+    def asDeterministic(self) -> "UserDefinedTableFunction":
+        """
+        Updates UserDefinedTableFunction to deterministic.
+        """
+        # Explicitly clean the cache to create a JVM UDTF instance.
+        self._judtf_placeholder = None
+        self.deterministic = True
+        return self
+
+
+class UDTFRegistration:
+    """
+    Wrapper for user-defined table function registration. This instance can be accessed by
+    :attr:`spark.udtf` or :attr:`sqlContext.udtf`.
+
+    .. versionadded:: 3.5.0
+    """
+
+    def __init__(self, sparkSession: "SparkSession"):
+        self.sparkSession = sparkSession
+
+    def register(
+        self,
+        name: str,
+        f: "UserDefinedTableFunction",
+    ) -> "UserDefinedTableFunction":
+        """Register a Python user-defined table function as a SQL table function.
+
+        .. versionadded:: 3.5.0
+
+        Parameters
+        ----------
+        name : str
+            The name of the user-defined table function in SQL statements.
+        f : function or :meth:`pyspark.sql.functions.udtf`
+            The user-defined table function.
+
+        Returns
+        -------
+        function
+            The registered user-defined table function.
+
+        Notes
+        -----
+        Spark uses the return type of the given user-defined table function as the return
+        type of the registered user-defined function.
+
+        To register a nondeterministic Python table function, users need to first build
+        a nondeterministic user-defined table function and then register it as a SQL function.
+
+        Examples
+        --------
+        >>> from pyspark.sql.functions import udtf
+        >>> @udtf(returnType="c1: int, c2: int")
+        ... class PlusOne:
+        ...     def eval(self, x: int):
+        ...         yield x, x + 1
+        ...
+        >>> _ = spark.udtf.register(name="plus_one", f=PlusOne)
+        >>> spark.sql("SELECT * FROM plus_one(1)").collect()
+        [Row(c1=1, c2=2)]
+
+        Use it with lateral join
+
+        >>> spark.sql("SELECT * FROM VALUES (0, 1), (1, 2) t(x, y), LATERAL plus_one(x)").collect()
+        [Row(x=0, y=1, c1=0, c2=1), Row(x=1, y=2, c1=1, c2=2)]
+        """
+        if f.evalType not in [PythonEvalType.SQL_TABLE_UDF, PythonEvalType.SQL_ARROW_TABLE_UDF]:
+            raise PySparkTypeError(
+                error_class="INVALID_UDTF_EVAL_TYPE",
+                message_parameters={
+                    "name": name,
+                    "eval_type": "SQL_TABLE_UDF, SQL_ARROW_TABLE_UDF",
+                },
+            )
+
+        register_udtf = _create_udtf(
+            cls=f.func,
+            returnType=f.returnType,
+            name=name,
+            evalType=f.evalType,
+            deterministic=f.deterministic,
+        )
+        self.sparkSession._jsparkSession.udtf().registerPython(name, register_udtf._judtf)
+        return register_udtf
+
+
+def _test() -> None:
+    import doctest
+    from pyspark.sql import SparkSession
+    import pyspark.sql.udf
+
+    globs = pyspark.sql.udtf.__dict__.copy()
+    spark = SparkSession.builder.master("local[4]").appName("sql.udtf tests").getOrCreate()
+    globs["spark"] = spark
+    (failure_count, test_count) = doctest.testmod(
+        pyspark.sql.udtf, globs=globs, optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE
+    )
+    spark.stop()
+    if failure_count:
+        sys.exit(-1)
+
+
+if __name__ == "__main__":
+    _test()

snowflake/snowpark_connect/includes/python/pyspark/sql/utils.py

@@ -0,0 +1,339 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import inspect
+import functools
+import os
+from typing import Any, Callable, Optional, Sequence, TYPE_CHECKING, cast, TypeVar, Union, Type
+
+from py4j.java_collections import JavaArray
+from py4j.java_gateway import (
+    JavaClass,
+    JavaGateway,
+    JavaObject,
+)
+
+from pyspark import SparkContext
+
+# For backward compatibility.
+from pyspark.errors import (  # noqa: F401
+    AnalysisException,
+    ParseException,
+    IllegalArgumentException,
+    StreamingQueryException,
+    QueryExecutionException,
+    PythonException,
+    UnknownException,
+    SparkUpgradeException,
+    PySparkNotImplementedError,
+)
+from pyspark.errors.exceptions.captured import CapturedException  # noqa: F401
+from pyspark.find_spark_home import _find_spark_home
+
+if TYPE_CHECKING:
+    from pyspark.sql.session import SparkSession
+    from pyspark.sql.dataframe import DataFrame
+    from pyspark.sql.column import Column
+    from pyspark.sql.window import Window
+    from pyspark.pandas._typing import IndexOpsLike, SeriesOrIndex
+
+has_numpy = False
+try:
+    import numpy as np  # noqa: F401
+
+    has_numpy = True
+except ImportError:
+    pass
+
+
+FuncT = TypeVar("FuncT", bound=Callable[..., Any])
+
+
+def toJArray(gateway: JavaGateway, jtype: JavaClass, arr: Sequence[Any]) -> JavaArray:
+    """
+    Convert python list to java type array
+
+    Parameters
+    ----------
+    gateway :
+        Py4j Gateway
+    jtype :
+        java type of element in array
+    arr :
+        python type list
+    """
+    jarray: JavaArray = gateway.new_array(jtype, len(arr))
+    for i in range(0, len(arr)):
+        jarray[i] = arr[i]
+    return jarray
+
+
+def require_test_compiled() -> None:
+    """Raise Exception if test classes are not compiled"""
+    import os
+    import glob
+
+    test_class_path = os.path.join(_find_spark_home(), "sql", "core", "target", "*", "test-classes")
+    paths = glob.glob(test_class_path)
+
+    if len(paths) == 0:
+        raise RuntimeError(
+            "%s doesn't exist. Spark sql test classes are not compiled." % test_class_path
+        )
+
+
+class ForeachBatchFunction:
+    """
+    This is the Python implementation of Java interface 'ForeachBatchFunction'. This wraps
+    the user-defined 'foreachBatch' function such that it can be called from the JVM when
+    the query is active.
+    """
+
+    def __init__(self, session: "SparkSession", func: Callable[["DataFrame", int], None]):
+        self.func = func
+        self.session = session
+
+    def call(self, jdf: JavaObject, batch_id: int) -> None:
+        from pyspark.sql.dataframe import DataFrame
+        from pyspark.sql.session import SparkSession
+
+        try:
+            session_jdf = jdf.sparkSession()
+            # assuming that spark context is still the same between JVM and PySpark
+            wrapped_session_jdf = SparkSession(self.session.sparkContext, session_jdf)
+            self.func(DataFrame(jdf, wrapped_session_jdf), batch_id)
+        except Exception as e:
+            self.error = e
+            raise e
+
+    class Java:
+        implements = ["org.apache.spark.sql.execution.streaming.sources.PythonForeachBatchFunction"]
+
+
+def to_str(value: Any) -> Optional[str]:
+    """
+    A wrapper over str(), but converts bool values to lower case strings.
+    If None is given, just returns None, instead of converting it to string "None".
+    """
+    if isinstance(value, bool):
+        return str(value).lower()
+    elif value is None:
+        return value
+    else:
+        return str(value)
+
+
+def is_timestamp_ntz_preferred() -> bool:
+    """
+    Return a bool if TimestampNTZType is preferred according to the SQL configuration set.
+    """
+    if is_remote():
+        from pyspark.sql.connect.session import SparkSession as ConnectSparkSession
+
+        session = ConnectSparkSession.getActiveSession()
+        if session is None:
+            return False
+        else:
+            return session.conf.get("spark.sql.timestampType", None) == "TIMESTAMP_NTZ"
+    else:
+        jvm = SparkContext._jvm
+        return jvm is not None and jvm.PythonSQLUtils.isTimestampNTZPreferred()
+
+
+def is_remote() -> bool:
+    """
+    Returns if the current running environment is for Spark Connect.
+    """
+    return "SPARK_CONNECT_MODE_ENABLED" in os.environ
+
+
+def try_remote_functions(f: FuncT) -> FuncT:
+    """Mark API supported from Spark Connect."""
+
+    @functools.wraps(f)
+    def wrapped(*args: Any, **kwargs: Any) -> Any:
+
+        if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
+            from pyspark.sql.connect import functions
+
+            return getattr(functions, f.__name__)(*args, **kwargs)
+        else:
+            return f(*args, **kwargs)
+
+    return cast(FuncT, wrapped)
+
+
+def try_remote_avro_functions(f: FuncT) -> FuncT:
+    """Mark API supported from Spark Connect."""
+
+    @functools.wraps(f)
+    def wrapped(*args: Any, **kwargs: Any) -> Any:
+
+        if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
+            from pyspark.sql.connect.avro import functions
+
+            return getattr(functions, f.__name__)(*args, **kwargs)
+        else:
+            return f(*args, **kwargs)
+
+    return cast(FuncT, wrapped)
+
+
+def try_remote_protobuf_functions(f: FuncT) -> FuncT:
+    """Mark API supported from Spark Connect."""
+
+    @functools.wraps(f)
+    def wrapped(*args: Any, **kwargs: Any) -> Any:
+
+        if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
+            from pyspark.sql.connect.protobuf import functions
+
+            return getattr(functions, f.__name__)(*args, **kwargs)
+        else:
+            return f(*args, **kwargs)
+
+    return cast(FuncT, wrapped)
+
+
+def try_remote_window(f: FuncT) -> FuncT:
+    """Mark API supported from Spark Connect."""
+
+    @functools.wraps(f)
+    def wrapped(*args: Any, **kwargs: Any) -> Any:
+
+        if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
+            from pyspark.sql.connect.window import Window  # type: ignore[misc]
+
+            return getattr(Window, f.__name__)(*args, **kwargs)
+        else:
+            return f(*args, **kwargs)
+
+    return cast(FuncT, wrapped)
+
+
+def try_remote_windowspec(f: FuncT) -> FuncT:
+    """Mark API supported from Spark Connect."""
+
+    @functools.wraps(f)
+    def wrapped(*args: Any, **kwargs: Any) -> Any:
+
+        if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
+            from pyspark.sql.connect.window import WindowSpec
+
+            return getattr(WindowSpec, f.__name__)(*args, **kwargs)
+        else:
+            return f(*args, **kwargs)
+
+    return cast(FuncT, wrapped)
+
+
+def get_active_spark_context() -> SparkContext:
+    """Raise RuntimeError if SparkContext is not initialized,
+    otherwise, returns the active SparkContext."""
+    sc = SparkContext._active_spark_context
+    if sc is None or sc._jvm is None:
+        raise RuntimeError("SparkContext or SparkSession should be created first.")
+    return sc
+
+
+def try_remote_observation(f: FuncT) -> FuncT:
+    """Mark API supported from Spark Connect."""
+
+    @functools.wraps(f)
+    def wrapped(*args: Any, **kwargs: Any) -> Any:
+        # TODO(SPARK-41527): Add the support of Observation.
+        if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
+            raise PySparkNotImplementedError(
+                error_class="NOT_IMPLEMENTED",
+                message_parameters={"feature": "Observation support for Spark Connect"},
+            )
+        return f(*args, **kwargs)
+
+    return cast(FuncT, wrapped)
+
+
+def try_remote_session_classmethod(f: FuncT) -> FuncT:
+    """Mark API supported from Spark Connect."""
+
+    @functools.wraps(f)
+    def wrapped(*args: Any, **kwargs: Any) -> Any:
+
+        if is_remote() and "PYSPARK_NO_NAMESPACE_SHARE" not in os.environ:
+            from pyspark.sql.connect.session import SparkSession  # type: ignore[misc]
+
+            assert inspect.isclass(args[0])
+            return getattr(SparkSession, f.__name__)(*args[1:], **kwargs)
+        else:
+            return f(*args, **kwargs)
+
+    return cast(FuncT, wrapped)
+
+
+def pyspark_column_op(
+    func_name: str, left: "IndexOpsLike", right: Any, fillna: Any = None
+) -> Union["SeriesOrIndex", None]:
+    """
+    Wrapper function for column_op to get proper Column class.
+    """
+    from pyspark.pandas.base import column_op
+    from pyspark.sql.column import Column as PySparkColumn
+    from pyspark.pandas.data_type_ops.base import _is_extension_dtypes
+
+    if is_remote():
+        from pyspark.sql.connect.column import Column as ConnectColumn
+
+        Column = ConnectColumn
+    else:
+        Column = PySparkColumn  # type: ignore[assignment]
+    result = column_op(getattr(Column, func_name))(left, right)
+    # It works as expected on extension dtype, so we don't need to call `fillna` for this case.
+    if (fillna is not None) and (_is_extension_dtypes(left) or _is_extension_dtypes(right)):
+        fillna = None
+    # TODO(SPARK-43877): Fix behavior difference for compare binary functions.
+    return result.fillna(fillna) if fillna is not None else result
+
+
+def get_column_class() -> Type["Column"]:
+    from pyspark.sql.column import Column as PySparkColumn
+
+    if is_remote():
+        from pyspark.sql.connect.column import Column as ConnectColumn
+
+        return ConnectColumn  # type: ignore[return-value]
+    else:
+        return PySparkColumn
+
+
+def get_dataframe_class() -> Type["DataFrame"]:
+    from pyspark.sql.dataframe import DataFrame as PySparkDataFrame
+
+    if is_remote():
+        from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame
+
+        return ConnectDataFrame  # type: ignore[return-value]
+    else:
+        return PySparkDataFrame
+
+
+def get_window_class() -> Type["Window"]:
+    from pyspark.sql.window import Window as PySparkWindow
+
+    if is_remote():
+        from pyspark.sql.connect.window import Window as ConnectWindow
+
+        return ConnectWindow  # type: ignore[return-value]
+    else:
+        return PySparkWindow