#
# 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
from typing import (
Any,
Generic,
List,
Optional,
TypeVar,
Union,
TYPE_CHECKING,
Tuple,
Callable,
)
import pandas as pd
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")
[docs]@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()
[docs] 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"
[docs]@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()
[docs] 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))
[docs] @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()
[docs]@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()