python/pyspark/testing/mlutils.py - spark - Git at Google

 #
 # 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 numpy as np

 from pyspark import keyword_only
 from pyspark.ml import Estimator, Model, Transformer, UnaryTransformer
 from pyspark.ml.evaluation import Evaluator
 from pyspark.ml.param import Param, Params, TypeConverters
 from pyspark.ml.param.shared import HasMaxIter, HasRegParam
 from pyspark.ml.classification import Classifier, ClassificationModel
 from pyspark.ml.util import DefaultParamsReadable, DefaultParamsWritable
 from pyspark.ml.wrapper import _java2py
 from pyspark.sql import DataFrame, SparkSession
 from pyspark.sql.types import DoubleType
 from pyspark.testing.utils import ReusedPySparkTestCase as PySparkTestCase


 def check_params(test_self, py_stage, check_params_exist=True):
     """
     Checks common requirements for :py:class:`PySpark.ml.Params.params`:

       - set of params exist in Java and Python and are ordered by names
       - param parent has the same UID as the object's UID
       - default param value from Java matches value in Python
       - optionally check if all params from Java also exist in Python
     """
     py_stage_str = "%s %s" % (type(py_stage), py_stage)
     if not hasattr(py_stage, "_to_java"):
         return
     java_stage = py_stage._to_java()
     if java_stage is None:
         return
     test_self.assertEqual(py_stage.uid, java_stage.uid(), msg=py_stage_str)
     if check_params_exist:
         param_names = [p.name for p in py_stage.params]
         java_params = list(java_stage.params())
         java_param_names = [jp.name() for jp in java_params]
         test_self.assertEqual(
             param_names,
             sorted(java_param_names),
             "Param list in Python does not match Java for %s:\nJava = %s\nPython = %s"
             % (py_stage_str, java_param_names, param_names),
         )
     for p in py_stage.params:
         test_self.assertEqual(p.parent, py_stage.uid)
         java_param = java_stage.getParam(p.name)
         py_has_default = py_stage.hasDefault(p)
         java_has_default = java_stage.hasDefault(java_param)
         test_self.assertEqual(
             py_has_default,
             java_has_default,
             "Default value mismatch of param %s for Params %s" % (p.name, str(py_stage)),
         )
         if py_has_default:
             if p.name == "seed":
                 continue  # Random seeds between Spark and PySpark are different
             java_default = _java2py(
                 test_self.sc, java_stage.clear(java_param).getOrDefault(java_param)
             )
             py_stage.clear(p)
             py_default = py_stage.getOrDefault(p)
             # equality test for NaN is always False
             if isinstance(java_default, float) and np.isnan(java_default):
                 java_default = "NaN"
                 py_default = "NaN" if np.isnan(py_default) else "not NaN"
             test_self.assertEqual(
                 java_default,
                 py_default,
                 "Java default %s != python default %s of param %s for Params %s"
                 % (str(java_default), str(py_default), p.name, str(py_stage)),
             )


 class SparkSessionTestCase(PySparkTestCase):
     @classmethod
     def setUpClass(cls):
         PySparkTestCase.setUpClass()
         cls.spark = SparkSession(cls.sc)

     @classmethod
     def tearDownClass(cls):
         PySparkTestCase.tearDownClass()
         cls.spark.stop()


 class MockDataset(DataFrame):
     def __new__(cls) -> "DataFrame":
         self = object.__new__(cls)
         self.__init__()
         return self

     def __init__(self):
         self.index = 0


 class HasFake(Params):
     def __init__(self):
         super(HasFake, self).__init__()
         self.fake = Param(self, "fake", "fake param")

     def getFake(self):
         return self.getOrDefault(self.fake)


 class MockTransformer(Transformer, HasFake):
     def __init__(self):
         super(MockTransformer, self).__init__()
         self.dataset_index = None

     def _transform(self, dataset):
         self.dataset_index = dataset.index
         dataset.index += 1
         return dataset


 class MockUnaryTransformer(UnaryTransformer, DefaultParamsReadable, DefaultParamsWritable):
     shift = Param(
         Params._dummy(),
         "shift",
         "The amount by which to shift " + "data in a DataFrame",
         typeConverter=TypeConverters.toFloat,
     )

     def __init__(self, shiftVal=1):
         super(MockUnaryTransformer, self).__init__()
         self._setDefault(shift=1)
         self._set(shift=shiftVal)

     def getShift(self):
         return self.getOrDefault(self.shift)

     def setShift(self, shift):
         self._set(shift=shift)

     def createTransformFunc(self):
         shiftVal = self.getShift()
         return lambda x: x + shiftVal

     def outputDataType(self):
         return DoubleType()

     def validateInputType(self, inputType):
         if inputType != DoubleType():
             raise TypeError("Bad input type: {}. ".format(inputType) + "Requires Double.")


 class MockEstimator(Estimator, HasFake):
     def __init__(self):
         super(MockEstimator, self).__init__()
         self.dataset_index = None

     def _fit(self, dataset):
         self.dataset_index = dataset.index
         model = MockModel()
         self._copyValues(model)
         return model


 class MockModel(MockTransformer, Model, HasFake):
     pass


 class _DummyLogisticRegressionParams(HasMaxIter, HasRegParam):
     def setMaxIter(self, value):
         return self._set(maxIter=value)

     def setRegParam(self, value):
         return self._set(regParam=value)


 # This is a dummy LogisticRegression used in test for python backend estimator/model
 class DummyLogisticRegression(
     Classifier, _DummyLogisticRegressionParams, DefaultParamsReadable, DefaultParamsWritable
 ):
     @keyword_only
     def __init__(
         self,
         *,
         featuresCol="features",
         labelCol="label",
         predictionCol="prediction",
         maxIter=100,
         regParam=0.0,
         rawPredictionCol="rawPrediction",
     ):
         super(DummyLogisticRegression, self).__init__()
         kwargs = self._input_kwargs
         self.setParams(**kwargs)

     @keyword_only
     def setParams(
         self,
         *,
         featuresCol="features",
         labelCol="label",
         predictionCol="prediction",
         maxIter=100,
         regParam=0.0,
         rawPredictionCol="rawPrediction",
     ):
         kwargs = self._input_kwargs
         self._set(**kwargs)
         return self

     def _fit(self, dataset):
         # Do nothing but create a dummy model
         return self._copyValues(DummyLogisticRegressionModel())


 class DummyLogisticRegressionModel(
     ClassificationModel,
     _DummyLogisticRegressionParams,
     DefaultParamsReadable,
     DefaultParamsWritable,
 ):
     def __init__(self):
         super(DummyLogisticRegressionModel, self).__init__()

     def _transform(self, dataset):
         # A dummy transform impl which always predict label 1
         from pyspark.sql.functions import array, lit
         from pyspark.ml.functions import array_to_vector

         rawPredCol = self.getRawPredictionCol()
         if rawPredCol:
             dataset = dataset.withColumn(
                 rawPredCol, array_to_vector(array(lit(-100.0), lit(100.0)))
             )
         predCol = self.getPredictionCol()
         if predCol:
             dataset = dataset.withColumn(predCol, lit(1.0))

         return dataset

     @property
     def numClasses(self):
         # a dummy implementation for test.
         return 2

     @property
     def intercept(self):
         # a dummy implementation for test.
         return 0.0

     # This class only used in test. The following methods/properties are not used in tests.

     @property
     def coefficients(self):
         raise NotImplementedError()

     def predictRaw(self, value):
         raise NotImplementedError()

     def numFeatures(self):
         raise NotImplementedError()

     def predict(self, value):
         raise NotImplementedError()


 class DummyEvaluator(Evaluator, DefaultParamsReadable, DefaultParamsWritable):
     def _evaluate(self, dataset):
         # a dummy implementation for test.
         return 1.0
	#
	# 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 numpy as np

	from pyspark import keyword_only
	from pyspark.ml import Estimator, Model, Transformer, UnaryTransformer
	from pyspark.ml.evaluation import Evaluator
	from pyspark.ml.param import Param, Params, TypeConverters
	from pyspark.ml.param.shared import HasMaxIter, HasRegParam
	from pyspark.ml.classification import Classifier, ClassificationModel
	from pyspark.ml.util import DefaultParamsReadable, DefaultParamsWritable
	from pyspark.ml.wrapper import _java2py
	from pyspark.sql import DataFrame, SparkSession
	from pyspark.sql.types import DoubleType
	from pyspark.testing.utils import ReusedPySparkTestCase as PySparkTestCase


	def check_params(test_self, py_stage, check_params_exist=True):
	"""
	Checks common requirements for :py:class:`PySpark.ml.Params.params`:

	- set of params exist in Java and Python and are ordered by names
	- param parent has the same UID as the object's UID
	- default param value from Java matches value in Python
	- optionally check if all params from Java also exist in Python
	"""
	py_stage_str = "%s %s" % (type(py_stage), py_stage)
	if not hasattr(py_stage, "_to_java"):
	return
	java_stage = py_stage._to_java()
	if java_stage is None:
	return
	test_self.assertEqual(py_stage.uid, java_stage.uid(), msg=py_stage_str)
	if check_params_exist:
	param_names = [p.name for p in py_stage.params]
	java_params = list(java_stage.params())
	java_param_names = [jp.name() for jp in java_params]
	test_self.assertEqual(
	param_names,
	sorted(java_param_names),
	"Param list in Python does not match Java for %s:\nJava = %s\nPython = %s"
	% (py_stage_str, java_param_names, param_names),
	)
	for p in py_stage.params:
	test_self.assertEqual(p.parent, py_stage.uid)
	java_param = java_stage.getParam(p.name)
	py_has_default = py_stage.hasDefault(p)
	java_has_default = java_stage.hasDefault(java_param)
	test_self.assertEqual(
	py_has_default,
	java_has_default,
	"Default value mismatch of param %s for Params %s" % (p.name, str(py_stage)),
	)
	if py_has_default:
	if p.name == "seed":
	continue # Random seeds between Spark and PySpark are different
	java_default = _java2py(
	test_self.sc, java_stage.clear(java_param).getOrDefault(java_param)
	)
	py_stage.clear(p)
	py_default = py_stage.getOrDefault(p)
	# equality test for NaN is always False
	if isinstance(java_default, float) and np.isnan(java_default):
	java_default = "NaN"
	py_default = "NaN" if np.isnan(py_default) else "not NaN"
	test_self.assertEqual(
	java_default,
	py_default,
	"Java default %s != python default %s of param %s for Params %s"
	% (str(java_default), str(py_default), p.name, str(py_stage)),
	)


	class SparkSessionTestCase(PySparkTestCase):
	@classmethod
	def setUpClass(cls):
	PySparkTestCase.setUpClass()
	cls.spark = SparkSession(cls.sc)

	@classmethod
	def tearDownClass(cls):
	PySparkTestCase.tearDownClass()
	cls.spark.stop()


	class MockDataset(DataFrame):
	def __new__(cls) -> "DataFrame":
	self = object.__new__(cls)
	self.__init__()
	return self

	def __init__(self):
	self.index = 0


	class HasFake(Params):
	def __init__(self):
	super(HasFake, self).__init__()
	self.fake = Param(self, "fake", "fake param")

	def getFake(self):
	return self.getOrDefault(self.fake)


	class MockTransformer(Transformer, HasFake):
	def __init__(self):
	super(MockTransformer, self).__init__()
	self.dataset_index = None

	def _transform(self, dataset):
	self.dataset_index = dataset.index
	dataset.index += 1
	return dataset


	class MockUnaryTransformer(UnaryTransformer, DefaultParamsReadable, DefaultParamsWritable):
	shift = Param(
	Params._dummy(),
	"shift",
	"The amount by which to shift " + "data in a DataFrame",
	typeConverter=TypeConverters.toFloat,
	)

	def __init__(self, shiftVal=1):
	super(MockUnaryTransformer, self).__init__()
	self._setDefault(shift=1)
	self._set(shift=shiftVal)

	def getShift(self):
	return self.getOrDefault(self.shift)

	def setShift(self, shift):
	self._set(shift=shift)

	def createTransformFunc(self):
	shiftVal = self.getShift()
	return lambda x: x + shiftVal

	def outputDataType(self):
	return DoubleType()

	def validateInputType(self, inputType):
	if inputType != DoubleType():
	raise TypeError("Bad input type: {}. ".format(inputType) + "Requires Double.")


	class MockEstimator(Estimator, HasFake):
	def __init__(self):
	super(MockEstimator, self).__init__()
	self.dataset_index = None

	def _fit(self, dataset):
	self.dataset_index = dataset.index
	model = MockModel()
	self._copyValues(model)
	return model


	class MockModel(MockTransformer, Model, HasFake):
	pass


	class _DummyLogisticRegressionParams(HasMaxIter, HasRegParam):
	def setMaxIter(self, value):
	return self._set(maxIter=value)

	def setRegParam(self, value):
	return self._set(regParam=value)


	# This is a dummy LogisticRegression used in test for python backend estimator/model
	class DummyLogisticRegression(
	Classifier, _DummyLogisticRegressionParams, DefaultParamsReadable, DefaultParamsWritable
	):
	@keyword_only
	def __init__(
	self,
	*,
	featuresCol="features",
	labelCol="label",
	predictionCol="prediction",
	maxIter=100,
	regParam=0.0,
	rawPredictionCol="rawPrediction",
	):
	super(DummyLogisticRegression, self).__init__()
	kwargs = self._input_kwargs
	self.setParams(**kwargs)

	@keyword_only
	def setParams(
	self,
	*,
	featuresCol="features",
	labelCol="label",
	predictionCol="prediction",
	maxIter=100,
	regParam=0.0,
	rawPredictionCol="rawPrediction",
	):
	kwargs = self._input_kwargs
	self._set(**kwargs)
	return self

	def _fit(self, dataset):
	# Do nothing but create a dummy model
	return self._copyValues(DummyLogisticRegressionModel())


	class DummyLogisticRegressionModel(
	ClassificationModel,
	_DummyLogisticRegressionParams,
	DefaultParamsReadable,
	DefaultParamsWritable,
	):
	def __init__(self):
	super(DummyLogisticRegressionModel, self).__init__()

	def _transform(self, dataset):
	# A dummy transform impl which always predict label 1
	from pyspark.sql.functions import array, lit
	from pyspark.ml.functions import array_to_vector

	rawPredCol = self.getRawPredictionCol()
	if rawPredCol:
	dataset = dataset.withColumn(
	rawPredCol, array_to_vector(array(lit(-100.0), lit(100.0)))
	)
	predCol = self.getPredictionCol()
	if predCol:
	dataset = dataset.withColumn(predCol, lit(1.0))

	return dataset

	@property
	def numClasses(self):
	# a dummy implementation for test.
	return 2

	@property
	def intercept(self):
	# a dummy implementation for test.
	return 0.0

	# This class only used in test. The following methods/properties are not used in tests.

	@property
	def coefficients(self):
	raise NotImplementedError()

	def predictRaw(self, value):
	raise NotImplementedError()

	def numFeatures(self):
	raise NotImplementedError()

	def predict(self, value):
	raise NotImplementedError()


	class DummyEvaluator(Evaluator, DefaultParamsReadable, DefaultParamsWritable):
	def _evaluate(self, dataset):
	# a dummy implementation for test.
	return 1.0