flink-ml-python/pyflink/ml/core/tests/test_stage.py - flink-ml - 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 os
 from typing import Dict, Any

 from pyflink.table import StreamTableEnvironment

 from pyflink.ml.core.api import Stage
 from pyflink.ml.core.linalg import Vectors
 from pyflink.ml.core.param import ParamValidators, Param, BooleanParam, IntParam, \
     FloatParam, StringParam, VectorParam, IntArrayParam, FloatArrayParam, StringArrayParam
 from pyflink.ml.tests.test_utils import PyFlinkMLTestCase

 BOOLEAN_PARAM = BooleanParam("boolean_param", "Description", False)
 INT_PARAM = IntParam("int_param", "Description", 1, ParamValidators.lt(100))
 FLOAT_PARAM = FloatParam("float_param", "Description", 3.0, ParamValidators.lt(100))
 STRING_PARAM = StringParam('string_param', "Description", "5")
 INT_ARRAY_PARAM = IntArrayParam("int_array_param", "Description", (6, 7))
 FLOAT_ARRAY_PARAM = FloatArrayParam("float_array_param", "Description", (10.0, 11.0))
 STRING_ARRAY_PARAM = StringArrayParam("string_array_param", "Description", ("14", "15"))
 VECTOR_PARAM = VectorParam('vector_param', "Description", Vectors.dense(1, 2, 3))
 EXTRA_INT_PARAM = IntParam("extra_int_param",
                            "Description",
                            20,
                            ParamValidators.always_true())
 PARAM_WITH_NONE_DEFAULT = IntParam("param_with_none_default",
                                    "Must be explicitly set with a non-none value",
                                    None,
                                    ParamValidators.not_null())


 class StageTest(PyFlinkMLTestCase):

     def test_param_set_value_with_name(self):
         stage = MyStage()
         stage.set(INT_PARAM, 2)
         self.assertEqual(2, stage.get(INT_PARAM))

         dense_vec = Vectors.dense(2, 2)
         stage.set(VECTOR_PARAM, dense_vec)
         self.assertEqual(dense_vec.get(0), stage.get(VECTOR_PARAM).get(0))
         self.assertEqual(dense_vec.get(1), stage.get(VECTOR_PARAM).get(1))

         sparse_vec = Vectors.sparse(3, [0, 2], [2, 2])
         stage.set(VECTOR_PARAM, sparse_vec)
         self.assertEqual(sparse_vec.get(0), stage.get(VECTOR_PARAM).get(0))
         self.assertEqual(sparse_vec.get(1), stage.get(VECTOR_PARAM).get(1))
         self.assertEqual(sparse_vec.get(2), stage.get(VECTOR_PARAM).get(2))

         param = stage.get_param("int_param")
         stage.set(param, 3)
         self.assertEqual(3, stage.get(param))

         param = stage.get_param('extra_int_param')
         stage.set(param, 50)
         self.assertEqual(50, stage.get(param))

     def test_param_with_null_default(self):
         stage = MyStage()
         import pytest
         with pytest.raises(ValueError, match='value should not be None'):
             stage.get(PARAM_WITH_NONE_DEFAULT)

         stage.set(PARAM_WITH_NONE_DEFAULT, 3)
         self.assertEqual(3, stage.get(PARAM_WITH_NONE_DEFAULT))

     def test_param_set_invalid_value(self):
         stage = MyStage()
         import pytest

         with pytest.raises(ValueError, match='Parameter int_param is given an invalid value 100.'):
             stage.set(INT_PARAM, 100)

         with pytest.raises(ValueError,
                            match='Parameter float_param is given an invalid value 100.0.'):
             stage.set(FLOAT_PARAM, 100.0)

         with pytest.raises(TypeError,
                            match="Parameter int_param's type <class 'int'> is incompatible with "
                                  "the type of <class 'str'>"):
             stage.set(INT_PARAM, "100")

         with pytest.raises(TypeError,
                            match="Parameter string_param's type <class 'str'> is incompatible with"
                                  " the type of <class 'int'>"):
             stage.set(STRING_PARAM, 100)

         with pytest.raises(TypeError,
                            match="Parameter vector_param's type <class 'pyflink.ml.core.linalg"
                                  ".Vector'> is incompatible with the type of <class 'int'>"):
             stage.set(VECTOR_PARAM, 100)

     def test_param_set_valid_value(self):
         stage = MyStage()

         stage.set(BOOLEAN_PARAM, True)
         self.assertTrue(stage.get(BOOLEAN_PARAM))

         stage.set(INT_PARAM, 50)
         self.assertEqual(50, stage.get(INT_PARAM))

         stage.set(FLOAT_PARAM, 50.0)
         self.assertEqual(50.0, stage.get(FLOAT_PARAM))

         stage.set(STRING_PARAM, "50")
         self.assertEqual("50", stage.get(STRING_PARAM))

         dense_vec = Vectors.dense(2, 2)
         stage.set(VECTOR_PARAM, dense_vec)
         self.assertEqual(dense_vec.get(0), stage.get(VECTOR_PARAM).get(0))
         self.assertEqual(dense_vec.get(1), stage.get(VECTOR_PARAM).get(1))

         sparse_vec = Vectors.sparse(3, [0, 2], [2, 2])
         stage.set(VECTOR_PARAM, sparse_vec)
         self.assertEqual(sparse_vec.get(0), stage.get(VECTOR_PARAM).get(0))
         self.assertEqual(sparse_vec.get(1), stage.get(VECTOR_PARAM).get(1))
         self.assertEqual(sparse_vec.get(2), stage.get(VECTOR_PARAM).get(2))

         stage.set(INT_ARRAY_PARAM, (50, 51))
         self.assertEqual((50, 51), stage.get(INT_ARRAY_PARAM))

         stage.set(FLOAT_ARRAY_PARAM, (50.0, 51.0))
         self.assertEqual((50.0, 51.0), stage.get(FLOAT_ARRAY_PARAM))

         stage.set(STRING_ARRAY_PARAM, ("50", "51"))
         self.assertEqual(("50", "51"), stage.get(STRING_ARRAY_PARAM))

     def test_stage_save_load(self):
         stage = MyStage()
         stage.set(PARAM_WITH_NONE_DEFAULT, 1)
         path = os.path.join(self.temp_dir, "test_stage_save_load")
         stage.save(path)
         loaded_stage = MyStage.load(self.env, path)
         self.assertEqual(stage.get_param_map(), loaded_stage.get_param_map())
         self.assertEqual(1, loaded_stage.get(INT_PARAM))

     def test_validators(self):
         gt = ParamValidators.gt(10)
         self.assertFalse(gt.validate(None))
         self.assertFalse(gt.validate(5))
         self.assertFalse(gt.validate(10))
         self.assertTrue(gt.validate(15))

         gt_eq = ParamValidators.gt_eq(10)
         self.assertFalse(gt_eq.validate(None))
         self.assertFalse(gt_eq.validate(5))
         self.assertTrue(gt_eq.validate(10))
         self.assertTrue(gt_eq.validate(15))

         lt = ParamValidators.lt(10)
         self.assertFalse(lt.validate(None))
         self.assertTrue(lt.validate(5))
         self.assertFalse(lt.validate(10))
         self.assertFalse(lt.validate(15))

         lt_eq = ParamValidators.lt_eq(10)
         self.assertFalse(lt_eq.validate(None))
         self.assertTrue(lt_eq.validate(5))
         self.assertTrue(lt_eq.validate(10))
         self.assertFalse(lt_eq.validate(15))

         in_range_inclusive = ParamValidators.in_range(5, 15)
         self.assertFalse(in_range_inclusive.validate(None))
         self.assertFalse(in_range_inclusive.validate(0))
         self.assertTrue(in_range_inclusive.validate(5))
         self.assertTrue(in_range_inclusive.validate(10))
         self.assertTrue(in_range_inclusive.validate(15))
         self.assertFalse(in_range_inclusive.validate(20))

         in_range_exclusive = ParamValidators.in_range(5, 15, False, False)
         self.assertFalse(in_range_exclusive.validate(None))
         self.assertFalse(in_range_exclusive.validate(0))
         self.assertFalse(in_range_exclusive.validate(5))
         self.assertTrue(in_range_exclusive.validate(10))
         self.assertFalse(in_range_exclusive.validate(15))
         self.assertFalse(in_range_exclusive.validate(20))

         in_array = ParamValidators.in_array([1, 2, 3])
         self.assertFalse(in_array.validate(None))
         self.assertTrue(in_array.validate(1))
         self.assertFalse(in_array.validate(0))

         not_null = ParamValidators.not_null()
         self.assertTrue(not_null.validate(5))
         self.assertFalse(not_null.validate(None))


 class MyStage(Stage):
     def __init__(self):
         self._param_map = {}  # type: Dict[Param, Any]
         self._init_param()

     def save(self, path: str) -> None:
         from pyflink.ml.util import read_write_utils
         read_write_utils.save_metadata(self, path)

     @classmethod
     def load(cls, t_env: StreamTableEnvironment, path: str):
         from pyflink.ml.util import read_write_utils
         return read_write_utils.load_stage_param(path)

     def get_param_map(self) -> Dict['Param[Any]', Any]:
         return self._param_map

     def _init_param(self):
         self._param_map[BOOLEAN_PARAM] = BOOLEAN_PARAM.default_value
         self._param_map[INT_PARAM] = INT_PARAM.default_value
         self._param_map[FLOAT_PARAM] = FLOAT_PARAM.default_value
         self._param_map[STRING_PARAM] = STRING_PARAM.default_value
         self._param_map[VECTOR_PARAM] = VECTOR_PARAM.default_value
         self._param_map[INT_ARRAY_PARAM] = INT_ARRAY_PARAM.default_value
         self._param_map[FLOAT_ARRAY_PARAM] = FLOAT_ARRAY_PARAM.default_value
         self._param_map[STRING_ARRAY_PARAM] = STRING_ARRAY_PARAM.default_value
         self._param_map[EXTRA_INT_PARAM] = EXTRA_INT_PARAM.default_value
         self._param_map[PARAM_WITH_NONE_DEFAULT] = PARAM_WITH_NONE_DEFAULT.default_value
	################################################################################
	# 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 os
	from typing import Dict, Any

	from pyflink.table import StreamTableEnvironment

	from pyflink.ml.core.api import Stage
	from pyflink.ml.core.linalg import Vectors
	from pyflink.ml.core.param import ParamValidators, Param, BooleanParam, IntParam, \
	FloatParam, StringParam, VectorParam, IntArrayParam, FloatArrayParam, StringArrayParam
	from pyflink.ml.tests.test_utils import PyFlinkMLTestCase

	BOOLEAN_PARAM = BooleanParam("boolean_param", "Description", False)
	INT_PARAM = IntParam("int_param", "Description", 1, ParamValidators.lt(100))
	FLOAT_PARAM = FloatParam("float_param", "Description", 3.0, ParamValidators.lt(100))
	STRING_PARAM = StringParam('string_param', "Description", "5")
	INT_ARRAY_PARAM = IntArrayParam("int_array_param", "Description", (6, 7))
	FLOAT_ARRAY_PARAM = FloatArrayParam("float_array_param", "Description", (10.0, 11.0))
	STRING_ARRAY_PARAM = StringArrayParam("string_array_param", "Description", ("14", "15"))
	VECTOR_PARAM = VectorParam('vector_param', "Description", Vectors.dense(1, 2, 3))
	EXTRA_INT_PARAM = IntParam("extra_int_param",
	"Description",
	20,
	ParamValidators.always_true())
	PARAM_WITH_NONE_DEFAULT = IntParam("param_with_none_default",
	"Must be explicitly set with a non-none value",
	None,
	ParamValidators.not_null())


	class StageTest(PyFlinkMLTestCase):

	def test_param_set_value_with_name(self):
	stage = MyStage()
	stage.set(INT_PARAM, 2)
	self.assertEqual(2, stage.get(INT_PARAM))

	dense_vec = Vectors.dense(2, 2)
	stage.set(VECTOR_PARAM, dense_vec)
	self.assertEqual(dense_vec.get(0), stage.get(VECTOR_PARAM).get(0))
	self.assertEqual(dense_vec.get(1), stage.get(VECTOR_PARAM).get(1))

	sparse_vec = Vectors.sparse(3, [0, 2], [2, 2])
	stage.set(VECTOR_PARAM, sparse_vec)
	self.assertEqual(sparse_vec.get(0), stage.get(VECTOR_PARAM).get(0))
	self.assertEqual(sparse_vec.get(1), stage.get(VECTOR_PARAM).get(1))
	self.assertEqual(sparse_vec.get(2), stage.get(VECTOR_PARAM).get(2))

	param = stage.get_param("int_param")
	stage.set(param, 3)
	self.assertEqual(3, stage.get(param))

	param = stage.get_param('extra_int_param')
	stage.set(param, 50)
	self.assertEqual(50, stage.get(param))

	def test_param_with_null_default(self):
	stage = MyStage()
	import pytest
	with pytest.raises(ValueError, match='value should not be None'):
	stage.get(PARAM_WITH_NONE_DEFAULT)

	stage.set(PARAM_WITH_NONE_DEFAULT, 3)
	self.assertEqual(3, stage.get(PARAM_WITH_NONE_DEFAULT))

	def test_param_set_invalid_value(self):
	stage = MyStage()
	import pytest

	with pytest.raises(ValueError, match='Parameter int_param is given an invalid value 100.'):
	stage.set(INT_PARAM, 100)

	with pytest.raises(ValueError,
	match='Parameter float_param is given an invalid value 100.0.'):
	stage.set(FLOAT_PARAM, 100.0)

	with pytest.raises(TypeError,
	match="Parameter int_param's type <class 'int'> is incompatible with "
	"the type of <class 'str'>"):
	stage.set(INT_PARAM, "100")

	with pytest.raises(TypeError,
	match="Parameter string_param's type <class 'str'> is incompatible with"
	" the type of <class 'int'>"):
	stage.set(STRING_PARAM, 100)

	with pytest.raises(TypeError,
	match="Parameter vector_param's type <class 'pyflink.ml.core.linalg"
	".Vector'> is incompatible with the type of <class 'int'>"):
	stage.set(VECTOR_PARAM, 100)

	def test_param_set_valid_value(self):
	stage = MyStage()

	stage.set(BOOLEAN_PARAM, True)
	self.assertTrue(stage.get(BOOLEAN_PARAM))

	stage.set(INT_PARAM, 50)
	self.assertEqual(50, stage.get(INT_PARAM))

	stage.set(FLOAT_PARAM, 50.0)
	self.assertEqual(50.0, stage.get(FLOAT_PARAM))

	stage.set(STRING_PARAM, "50")
	self.assertEqual("50", stage.get(STRING_PARAM))

	dense_vec = Vectors.dense(2, 2)
	stage.set(VECTOR_PARAM, dense_vec)
	self.assertEqual(dense_vec.get(0), stage.get(VECTOR_PARAM).get(0))
	self.assertEqual(dense_vec.get(1), stage.get(VECTOR_PARAM).get(1))

	sparse_vec = Vectors.sparse(3, [0, 2], [2, 2])
	stage.set(VECTOR_PARAM, sparse_vec)
	self.assertEqual(sparse_vec.get(0), stage.get(VECTOR_PARAM).get(0))
	self.assertEqual(sparse_vec.get(1), stage.get(VECTOR_PARAM).get(1))
	self.assertEqual(sparse_vec.get(2), stage.get(VECTOR_PARAM).get(2))

	stage.set(INT_ARRAY_PARAM, (50, 51))
	self.assertEqual((50, 51), stage.get(INT_ARRAY_PARAM))

	stage.set(FLOAT_ARRAY_PARAM, (50.0, 51.0))
	self.assertEqual((50.0, 51.0), stage.get(FLOAT_ARRAY_PARAM))

	stage.set(STRING_ARRAY_PARAM, ("50", "51"))
	self.assertEqual(("50", "51"), stage.get(STRING_ARRAY_PARAM))

	def test_stage_save_load(self):
	stage = MyStage()
	stage.set(PARAM_WITH_NONE_DEFAULT, 1)
	path = os.path.join(self.temp_dir, "test_stage_save_load")
	stage.save(path)
	loaded_stage = MyStage.load(self.env, path)
	self.assertEqual(stage.get_param_map(), loaded_stage.get_param_map())
	self.assertEqual(1, loaded_stage.get(INT_PARAM))

	def test_validators(self):
	gt = ParamValidators.gt(10)
	self.assertFalse(gt.validate(None))
	self.assertFalse(gt.validate(5))
	self.assertFalse(gt.validate(10))
	self.assertTrue(gt.validate(15))

	gt_eq = ParamValidators.gt_eq(10)
	self.assertFalse(gt_eq.validate(None))
	self.assertFalse(gt_eq.validate(5))
	self.assertTrue(gt_eq.validate(10))
	self.assertTrue(gt_eq.validate(15))

	lt = ParamValidators.lt(10)
	self.assertFalse(lt.validate(None))
	self.assertTrue(lt.validate(5))
	self.assertFalse(lt.validate(10))
	self.assertFalse(lt.validate(15))

	lt_eq = ParamValidators.lt_eq(10)
	self.assertFalse(lt_eq.validate(None))
	self.assertTrue(lt_eq.validate(5))
	self.assertTrue(lt_eq.validate(10))
	self.assertFalse(lt_eq.validate(15))

	in_range_inclusive = ParamValidators.in_range(5, 15)
	self.assertFalse(in_range_inclusive.validate(None))
	self.assertFalse(in_range_inclusive.validate(0))
	self.assertTrue(in_range_inclusive.validate(5))
	self.assertTrue(in_range_inclusive.validate(10))
	self.assertTrue(in_range_inclusive.validate(15))
	self.assertFalse(in_range_inclusive.validate(20))

	in_range_exclusive = ParamValidators.in_range(5, 15, False, False)
	self.assertFalse(in_range_exclusive.validate(None))
	self.assertFalse(in_range_exclusive.validate(0))
	self.assertFalse(in_range_exclusive.validate(5))
	self.assertTrue(in_range_exclusive.validate(10))
	self.assertFalse(in_range_exclusive.validate(15))
	self.assertFalse(in_range_exclusive.validate(20))

	in_array = ParamValidators.in_array([1, 2, 3])
	self.assertFalse(in_array.validate(None))
	self.assertTrue(in_array.validate(1))
	self.assertFalse(in_array.validate(0))

	not_null = ParamValidators.not_null()
	self.assertTrue(not_null.validate(5))
	self.assertFalse(not_null.validate(None))


	class MyStage(Stage):
	def __init__(self):
	self._param_map = {} # type: Dict[Param, Any]
	self._init_param()

	def save(self, path: str) -> None:
	from pyflink.ml.util import read_write_utils
	read_write_utils.save_metadata(self, path)

	@classmethod
	def load(cls, t_env: StreamTableEnvironment, path: str):
	from pyflink.ml.util import read_write_utils
	return read_write_utils.load_stage_param(path)

	def get_param_map(self) -> Dict['Param[Any]', Any]:
	return self._param_map

	def _init_param(self):
	self._param_map[BOOLEAN_PARAM] = BOOLEAN_PARAM.default_value
	self._param_map[INT_PARAM] = INT_PARAM.default_value
	self._param_map[FLOAT_PARAM] = FLOAT_PARAM.default_value
	self._param_map[STRING_PARAM] = STRING_PARAM.default_value
	self._param_map[VECTOR_PARAM] = VECTOR_PARAM.default_value
	self._param_map[INT_ARRAY_PARAM] = INT_ARRAY_PARAM.default_value
	self._param_map[FLOAT_ARRAY_PARAM] = FLOAT_ARRAY_PARAM.default_value
	self._param_map[STRING_ARRAY_PARAM] = STRING_ARRAY_PARAM.default_value
	self._param_map[EXTRA_INT_PARAM] = EXTRA_INT_PARAM.default_value
	self._param_map[PARAM_WITH_NONE_DEFAULT] = PARAM_WITH_NONE_DEFAULT.default_value