snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/includes/python/pyspark/ml/common.py

@@ -0,0 +1,146 @@
+#
+# 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.
+#
+
+from typing import Any, Callable, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pyspark.ml._typing import C, JavaObjectOrPickleDump
+
+import py4j.protocol
+from py4j.protocol import Py4JJavaError
+from py4j.java_gateway import JavaObject
+from py4j.java_collections import JavaArray, JavaList
+
+import pyspark.context
+from pyspark import RDD, SparkContext
+from pyspark.serializers import CPickleSerializer, AutoBatchedSerializer
+from pyspark.sql import DataFrame, SparkSession
+
+# Hack for support float('inf') in Py4j
+_old_smart_decode = py4j.protocol.smart_decode
+
+_float_str_mapping = {
+    "nan": "NaN",
+    "inf": "Infinity",
+    "-inf": "-Infinity",
+}
+
+
+def _new_smart_decode(obj: Any) -> str:
+    if isinstance(obj, float):
+        s = str(obj)
+        return _float_str_mapping.get(s, s)
+    return _old_smart_decode(obj)
+
+
+py4j.protocol.smart_decode = _new_smart_decode
+
+
+_picklable_classes = [
+    "SparseVector",
+    "DenseVector",
+    "SparseMatrix",
+    "DenseMatrix",
+]
+
+
+# this will call the ML version of pythonToJava()
+def _to_java_object_rdd(rdd: RDD) -> JavaObject:
+    """Return an JavaRDD of Object by unpickling
+
+    It will convert each Python object into Java object by Pickle, whenever the
+    RDD is serialized in batch or not.
+    """
+    rdd = rdd._reserialize(AutoBatchedSerializer(CPickleSerializer()))
+    assert rdd.ctx._jvm is not None
+    return rdd.ctx._jvm.org.apache.spark.ml.python.MLSerDe.pythonToJava(rdd._jrdd, True)
+
+
+def _py2java(sc: SparkContext, obj: Any) -> JavaObject:
+    """Convert Python object into Java"""
+    if isinstance(obj, RDD):
+        obj = _to_java_object_rdd(obj)
+    elif isinstance(obj, DataFrame):
+        obj = obj._jdf
+    elif isinstance(obj, SparkContext):
+        obj = obj._jsc
+    elif isinstance(obj, list):
+        obj = [_py2java(sc, x) for x in obj]
+    elif isinstance(obj, JavaObject):
+        pass
+    elif isinstance(obj, (int, float, bool, bytes, str)):
+        pass
+    else:
+        data = bytearray(CPickleSerializer().dumps(obj))
+        assert sc._jvm is not None
+        obj = sc._jvm.org.apache.spark.ml.python.MLSerDe.loads(data)
+    return obj
+
+
+def _java2py(sc: SparkContext, r: "JavaObjectOrPickleDump", encoding: str = "bytes") -> Any:
+    if isinstance(r, JavaObject):
+        clsName = r.getClass().getSimpleName()
+        # convert RDD into JavaRDD
+        if clsName != "JavaRDD" and clsName.endswith("RDD"):
+            r = r.toJavaRDD()
+            clsName = "JavaRDD"
+
+        assert sc._jvm is not None
+
+        if clsName == "JavaRDD":
+            jrdd = sc._jvm.org.apache.spark.ml.python.MLSerDe.javaToPython(r)
+            return RDD(jrdd, sc)
+
+        if clsName == "Dataset":
+            return DataFrame(r, SparkSession._getActiveSessionOrCreate())
+
+        if clsName in _picklable_classes:
+            r = sc._jvm.org.apache.spark.ml.python.MLSerDe.dumps(r)
+        elif isinstance(r, (JavaArray, JavaList)):
+            try:
+                r = sc._jvm.org.apache.spark.ml.python.MLSerDe.dumps(r)
+            except Py4JJavaError:
+                pass  # not picklable
+
+    if isinstance(r, (bytearray, bytes)):
+        r = CPickleSerializer().loads(bytes(r), encoding=encoding)
+    return r
+
+
+def callJavaFunc(
+    sc: pyspark.context.SparkContext, func: Callable[..., "JavaObjectOrPickleDump"], *args: Any
+) -> "JavaObjectOrPickleDump":
+    """Call Java Function"""
+    java_args = [_py2java(sc, a) for a in args]
+    return _java2py(sc, func(*java_args))
+
+
+def inherit_doc(cls: "C") -> "C":
+    """
+    A decorator that makes a class inherit documentation from its parents.
+    """
+    for name, func in vars(cls).items():
+        # only inherit docstring for public functions
+        if name.startswith("_"):
+            continue
+        if not func.__doc__:
+            for parent in cls.__bases__:
+                parent_func = getattr(parent, name, None)
+                if parent_func and getattr(parent_func, "__doc__", None):
+                    func.__doc__ = parent_func.__doc__
+                    break
+    return cls

snowflake/snowpark_connect/includes/python/pyspark/ml/connect/__init__.py

@@ -0,0 +1,44 @@
+#
+# 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.
+#
+
+"""Spark Connect Python Client - ML module"""
+
+from pyspark.ml.connect.base import (
+    Estimator,
+    Transformer,
+    Model,
+)
+
+from pyspark.ml.connect import (
+    feature,
+    evaluation,
+    tuning,
+)
+
+from pyspark.ml.connect.pipeline import Pipeline, PipelineModel
+
+__all__ = [
+    "Estimator",
+    "Transformer",
+    "Estimator",
+    "Model",
+    "feature",
+    "evaluation",
+    "Pipeline",
+    "PipelineModel",
+    "tuning",
+]

snowflake/snowpark_connect/includes/python/pyspark/ml/connect/base.py

@@ -0,0 +1,346 @@
+#
+# 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.
+#
+
+from abc import ABCMeta, abstractmethod
+
+import pandas as pd
+
+from typing import (
+    Any,
+    Generic,
+    List,
+    Optional,
+    TypeVar,
+    Union,
+    TYPE_CHECKING,
+    Tuple,
+    Callable,
+)
+
+from pyspark import since
+from pyspark.ml.common import inherit_doc
+from pyspark.sql.dataframe import DataFrame
+from pyspark.ml.param import Params
+from pyspark.ml.param.shared import (
+    HasLabelCol,
+    HasFeaturesCol,
+    HasPredictionCol,
+)
+from pyspark.ml.connect.util import transform_dataframe_column
+
+if TYPE_CHECKING:
+    from pyspark.ml._typing import ParamMap
+
+M = TypeVar("M", bound="Transformer")
+
+
+@inherit_doc
+class Estimator(Params, Generic[M], metaclass=ABCMeta):
+    """
+    Abstract class for estimators that fit models to data.
+
+    .. versionadded:: 3.5.0
+    """
+
+    @abstractmethod
+    def _fit(self, dataset: Union[DataFrame, pd.DataFrame]) -> M:
+        """
+        Fits a model to the input dataset. This is called by the default implementation of fit.
+
+
+        Parameters
+        ----------
+        dataset : :py:class:`pyspark.sql.DataFrame`
+            input dataset
+
+        Returns
+        -------
+        :class:`Transformer`
+            fitted model
+        """
+        raise NotImplementedError()
+
+    def fit(
+        self,
+        dataset: Union[DataFrame, pd.DataFrame],
+        params: Optional["ParamMap"] = None,
+    ) -> Union[M, List[M]]:
+        """
+        Fits a model to the input dataset with optional parameters.
+
+        .. versionadded:: 3.5.0
+
+        Parameters
+        ----------
+        dataset : :py:class:`pyspark.sql.DataFrame` or py:class:`pandas.DataFrame`
+            input dataset, it can be either pandas dataframe or spark dataframe.
+        params : a dict of param values, optional
+            an optional param map that overrides embedded params.
+
+        Returns
+        -------
+        :py:class:`Transformer`
+            fitted model
+        """
+        if params is None:
+            params = dict()
+
+        if isinstance(params, dict):
+            if params:
+                return self.copy(params)._fit(dataset)
+            else:
+                return self._fit(dataset)
+        else:
+            raise TypeError(
+                "Params must be either a param map or a list/tuple of param maps, "
+                "but got %s." % type(params)
+            )
+
+
+_SPARKML_TRANSFORMER_TMP_OUTPUT_COLNAME = "_sparkML_transformer_tmp_output"
+
+
+@inherit_doc
+class Transformer(Params, metaclass=ABCMeta):
+    """
+    Abstract class for transformers that transform one dataset into another.
+
+    .. versionadded:: 3.5.0
+    """
+
+    def _input_columns(self) -> List[str]:
+        """
+        Return a list of input column names which are used as inputs of transformation.
+        """
+        raise NotImplementedError()
+
+    def _output_columns(self) -> List[Tuple[str, str]]:
+        """
+        Return a list of output transformed columns, each elements in the list
+        is a tuple of (column_name, column_spark_type)
+        """
+        raise NotImplementedError()
+
+    def _get_transform_fn(self) -> Callable[["pd.Series"], Any]:
+        """
+        Return a transformation function that accepts an instance of `pd.Series` as input and
+        returns transformed result as an instance of `pd.Series` or `pd.DataFrame`.
+        If there's only one output column, the transformed result must be an
+        instance of `pd.Series`, if there are multiple output columns, the transformed result
+        must be an instance of `pd.DataFrame` with column names matching output schema
+        returned by  `_output_columns` interface.
+        """
+        raise NotImplementedError()
+
+    def transform(
+        self, dataset: Union[DataFrame, pd.DataFrame], params: Optional["ParamMap"] = None
+    ) -> Union[DataFrame, pd.DataFrame]:
+        """
+        Transforms the input dataset.
+        The dataset can be either pandas dataframe or spark dataframe,
+        if it is a spark DataFrame, the result of transformation is a new spark DataFrame
+        that contains all existing columns and output columns with names.
+        if it is a pandas DataFrame, the input pandas dataframe is appended with output
+        columns in place.
+
+        Note: Transformers does not allow output column having the same name with
+        existing columns.
+
+        Parameters
+        ----------
+        dataset : :py:class:`pyspark.sql.DataFrame` or py:class:`pandas.DataFrame`
+            input dataset.
+
+        params : dict, optional
+            an optional param map that overrides embedded params.
+
+        Returns
+        -------
+        :py:class:`pyspark.sql.DataFrame` or py:class:`pandas.DataFrame`
+            transformed dataset, the type of output dataframe is consistent with
+            input dataframe.
+        """
+        if params is None:
+            params = dict()
+        if isinstance(params, dict):
+            if params:
+                return self.copy(params)._transform(dataset)
+            else:
+                return self._transform(dataset)
+
+    def _transform(self, dataset: Union[DataFrame, pd.DataFrame]) -> Union[DataFrame, pd.DataFrame]:
+        input_cols = self._input_columns()
+        transform_fn = self._get_transform_fn()
+        output_cols = self._output_columns()
+
+        existing_cols = list(dataset.columns)
+        for col_name, _ in output_cols:
+            if col_name in existing_cols:
+                raise ValueError(
+                    "Transformers does not allow output column having the same name with "
+                    "existing columns."
+                )
+
+        return transform_dataframe_column(
+            dataset,
+            input_cols=input_cols,
+            transform_fn=transform_fn,
+            output_cols=output_cols,
+        )
+
+
+@inherit_doc
+class Evaluator(Params, metaclass=ABCMeta):
+    """
+    Base class for evaluators that compute metrics from predictions.
+
+    .. versionadded:: 3.5.0
+    """
+
+    @abstractmethod
+    def _evaluate(self, dataset: Union["DataFrame", "pd.DataFrame"]) -> float:
+        """
+        Evaluates the output.
+
+        Parameters
+        ----------
+        dataset : :py:class:`pyspark.sql.DataFrame`
+            a dataset that contains labels/observations and predictions
+
+        Returns
+        -------
+        float
+            metric
+        """
+        raise NotImplementedError()
+
+    def evaluate(self, dataset: DataFrame, params: Optional["ParamMap"] = None) -> float:
+        """
+        Evaluates the output with optional parameters.
+
+        .. versionadded:: 3.5.0
+
+        Parameters
+        ----------
+        dataset : :py:class:`pyspark.sql.DataFrame`
+            a dataset that contains labels/observations and predictions
+        params : dict, optional
+            an optional param map that overrides embedded params
+
+        Returns
+        -------
+        float
+            metric
+        """
+        if params is None:
+            params = dict()
+        if isinstance(params, dict):
+            if params:
+                return self.copy(params)._evaluate(dataset)
+            else:
+                return self._evaluate(dataset)
+        else:
+            raise TypeError("Params must be a param map but got %s." % type(params))
+
+    @since("1.5.0")
+    def isLargerBetter(self) -> bool:
+        """
+        Indicates whether the metric returned by :py:meth:`evaluate` should be maximized
+        (True, default) or minimized (False).
+        A given evaluator may support multiple metrics which may be maximized or minimized.
+        """
+        raise NotImplementedError()
+
+
+@inherit_doc
+class Model(Transformer, metaclass=ABCMeta):
+    """
+    Abstract class for models that are fitted by estimators.
+
+    .. versionadded:: 3.5.0
+    """
+
+    pass
+
+
+@inherit_doc
+class _PredictorParams(HasLabelCol, HasFeaturesCol, HasPredictionCol):
+    """
+    Params for :py:class:`Predictor` and :py:class:`PredictorModel`.
+
+    .. versionadded:: 3.5.0
+    """
+
+    pass
+
+
+@inherit_doc
+class Predictor(Estimator[M], _PredictorParams, metaclass=ABCMeta):
+    """
+    Estimator for prediction tasks (regression and classification).
+    """
+
+    @since("3.5.0")
+    def setLabelCol(self, value: str) -> "Predictor":
+        """
+        Sets the value of :py:attr:`labelCol`.
+        """
+        return self._set(labelCol=value)
+
+    @since("3.5.0")
+    def setFeaturesCol(self, value: str) -> "Predictor":
+        """
+        Sets the value of :py:attr:`featuresCol`.
+        """
+        return self._set(featuresCol=value)
+
+    @since("3.5.0")
+    def setPredictionCol(self, value: str) -> "Predictor":
+        """
+        Sets the value of :py:attr:`predictionCol`.
+        """
+        return self._set(predictionCol=value)
+
+
+@inherit_doc
+class PredictionModel(Model, _PredictorParams, metaclass=ABCMeta):
+    """
+    Model for prediction tasks (regression and classification).
+    """
+
+    @since("3.5.0")
+    def setFeaturesCol(self, value: str) -> "PredictionModel":
+        """
+        Sets the value of :py:attr:`featuresCol`.
+        """
+        return self._set(featuresCol=value)
+
+    @since("3.5.0")
+    def setPredictionCol(self, value: str) -> "PredictionModel":
+        """
+        Sets the value of :py:attr:`predictionCol`.
+        """
+        return self._set(predictionCol=value)
+
+    @property
+    @abstractmethod
+    @since("3.5.0")
+    def numFeatures(self) -> int:
+        """
+        Returns the number of features the model was trained on. If unknown, returns -1
+        """
+        raise NotImplementedError()