flink-ml-core/src/test/java/org/apache/flink/ml/api/StageTest.java - 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.
  */

 package org.apache.flink.ml.api;

 import org.apache.flink.ml.linalg.Vector;
 import org.apache.flink.ml.linalg.Vectors;
 import org.apache.flink.ml.param.BooleanParam;
 import org.apache.flink.ml.param.DoubleArrayArrayParam;
 import org.apache.flink.ml.param.DoubleArrayParam;
 import org.apache.flink.ml.param.DoubleParam;
 import org.apache.flink.ml.param.FloatArrayParam;
 import org.apache.flink.ml.param.FloatParam;
 import org.apache.flink.ml.param.IntArrayParam;
 import org.apache.flink.ml.param.IntParam;
 import org.apache.flink.ml.param.LongArrayParam;
 import org.apache.flink.ml.param.LongParam;
 import org.apache.flink.ml.param.Param;
 import org.apache.flink.ml.param.ParamValidator;
 import org.apache.flink.ml.param.ParamValidators;
 import org.apache.flink.ml.param.StringArrayArrayParam;
 import org.apache.flink.ml.param.StringArrayParam;
 import org.apache.flink.ml.param.StringParam;
 import org.apache.flink.ml.param.VectorParam;
 import org.apache.flink.ml.param.WithParams;
 import org.apache.flink.ml.util.ParamUtils;
 import org.apache.flink.ml.util.ReadWriteUtils;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
 import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;

 import org.junit.Assert;
 import org.junit.Test;

 import java.io.IOException;
 import java.nio.file.Files;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;

 /** Tests the behavior of Stage and WithParams. */
 public class StageTest {

     // A WithParams subclass which has one parameter for each pre-defined parameter type.
     private interface MyParams<T> extends WithParams<T> {
         Param<Boolean> BOOLEAN_PARAM = new BooleanParam("booleanParam", "Description", false);

         Param<Integer> INT_PARAM =
                 new IntParam("intParam", "Description", 1, ParamValidators.lt(100));

         // This param might be regarded as integer type given a small value.
         Param<Long> LONG_PARAM =
                 new LongParam("longParam", "Description", 2L, ParamValidators.lt(100));

         // This param must be regarded as long type given a large enough value.
         Param<Long> LONG_PARAM2 =
                 new LongParam(
                         "longParam2",
                         "Description",
                         Integer.MAX_VALUE + 1L,
                         ParamValidators.lt(Integer.MAX_VALUE + 100L));

         Param<Float> FLOAT_PARAM =
                 new FloatParam("floatParam", "Description", 3.0f, ParamValidators.lt(100));

         Param<Double> DOUBLE_PARAM =
                 new DoubleParam("doubleParam", "Description", 4.0, ParamValidators.lt(100));

         Param<String> STRING_PARAM = new StringParam("stringParam", "Description", "5");

         Param<Integer[]> INT_ARRAY_PARAM =
                 new IntArrayParam("intArrayParam", "Description", new Integer[] {6, 7});

         Param<Long[]> LONG_ARRAY_PARAM =
                 new LongArrayParam(
                         "longArrayParam",
                         "Description",
                         new Long[] {8L, 9L},
                         ParamValidators.alwaysTrue());

         Param<Float[]> FLOAT_ARRAY_PARAM =
                 new FloatArrayParam("floatArrayParam", "Description", new Float[] {10.0f, 11.0f});

         Param<Double[]> DOUBLE_ARRAY_PARAM =
                 new DoubleArrayParam(
                         "doubleArrayParam",
                         "Description",
                         new Double[] {12.0, 13.0},
                         ParamValidators.alwaysTrue());

         Param<String[]> STRING_ARRAY_PARAM =
                 new StringArrayParam("stringArrayParam", "Description", new String[] {"14", "15"});

         Param<String[][]> STRING_ARRAY_ARRAY_PARAM =
                 new StringArrayArrayParam(
                         "stringArrayArrayParam",
                         "Description",
                         new String[][] {new String[] {"14", "15"}});

         Param<Double[][]> DOUBLE_ARRAY_ARRAY_PARAM =
                 new DoubleArrayArrayParam(
                         "doubleArrayArrayParam",
                         "Description",
                         new Double[][] {new Double[] {14.0, 15.0}, new Double[] {16.0, 17.0}});

         Param<Vector> VECTOR_PARAM =
                 new VectorParam("vectorParam", "Description", Vectors.dense(1.0, 2.0, 3.0));
     }

     /**
      * A Stage subclass which inherits all parameters from MyParams and defines an extra parameter.
      */
     public static class MyStage implements Stage<MyStage>, MyParams<MyStage> {
         private final Map<Param<?>, Object> paramMap = new HashMap<>();

         public final Param<Integer> extraIntParam =
                 new IntParam("extraIntParam", "Description", 20, ParamValidators.alwaysTrue());

         public final Param<Integer> paramWithNullDefault =
                 new IntParam(
                         "paramWithNullDefault",
                         "Must be explicitly set with a non-null value",
                         null,
                         ParamValidators.notNull());

         public MyStage() {
             ParamUtils.initializeMapWithDefaultValues(paramMap, this);
         }

         @Override
         public Map<Param<?>, Object> getParamMap() {
             return paramMap;
         }

         @Override
         public void save(String path) throws IOException {
             ReadWriteUtils.saveMetadata(this, path);
         }

         public static MyStage load(StreamTableEnvironment tEnv, String path) throws IOException {
             return ReadWriteUtils.loadStageParam(path);
         }
     }

     /** A Stage subclass without the static load() method. */
     public static class MyStageWithoutLoad implements Stage<MyStage>, MyParams<MyStage> {
         private final Map<Param<?>, Object> paramMap = new HashMap<>();

         public MyStageWithoutLoad() {
             ParamUtils.initializeMapWithDefaultValues(paramMap, this);
         }

         @Override
         public void save(String path) throws IOException {
             ReadWriteUtils.saveMetadata(this, path);
         }

         @Override
         public Map<Param<?>, Object> getParamMap() {
             return paramMap;
         }
     }

     // Asserts that m1 and m2 are equivalent.
     private static void assertParamMapEquals(Map<Param<?>, Object> m1, Map<Param<?>, Object> m2) {
         Assert.assertTrue(m1 != null && m2 != null);
         Assert.assertEquals(m1.size(), m2.size());

         for (Map.Entry<Param<?>, Object> entry : m1.entrySet()) {
             Assert.assertTrue(m2.containsKey(entry.getKey()));
             Object v1 = entry.getValue();
             Object v2 = m2.get(entry.getKey());
             if (v1 == null || v2 == null) {
                 Assert.assertTrue(v1 == null && v2 == null);
             } else if (v1.getClass().isArray() && v2.getClass().isArray()) {
                 Assert.assertArrayEquals((Object[]) v1, (Object[]) v2);
             } else {
                 Assert.assertEquals(v1, v2);
             }
         }
     }

     // Saves and loads the given stage. And verifies that the loaded stage has same parameter values
     // as the original stage.
     private static Stage<?> validateStageSaveLoad(
             StreamTableEnvironment tEnv, Stage<?> stage, Map<String, Object> paramOverrides)
             throws IOException {
         for (Map.Entry<String, Object> entry : paramOverrides.entrySet()) {
             Param<?> param = stage.getParam(entry.getKey());
             ReadWriteUtils.setParam(stage, param, entry.getValue());
         }

         String path = Files.createTempDirectory("").toString();
         stage.save(path);
         try {
             stage.save(path);
             Assert.fail("Expected IOException");
         } catch (IOException e) {
             // This is expected.
         }

         Stage<?> loadedStage = ReadWriteUtils.loadStage(tEnv, path);
         for (Map.Entry<String, Object> entry : paramOverrides.entrySet()) {
             Param<?> param = loadedStage.getParam(entry.getKey());
             Assert.assertEquals(entry.getValue(), loadedStage.get(param));
         }
         assertParamMapEquals(stage.getParamMap(), loadedStage.getParamMap());
         return loadedStage;
     }

     @Test
     public void testParamSetValueWithName() {
         MyStage stage = new MyStage();

         Param<Integer> paramA = MyParams.INT_PARAM;
         stage.set(paramA, 2);
         Assert.assertEquals(2, (int) stage.get(paramA));

         Param<Integer> paramB = stage.getParam("intParam");
         stage.set(paramB, 3);
         Assert.assertEquals(3, (int) stage.get(paramB));

         Param<Integer> paramC = stage.getParam("extraIntParam");
         stage.set(paramC, 50);
         Assert.assertEquals(50, (int) stage.get(paramC));
     }

     @Test
     public void testParamWithNullDefault() {
         MyStage stage = new MyStage();
         try {
             stage.get(stage.paramWithNullDefault);
             Assert.fail("Expected IllegalArgumentException");
         } catch (IllegalArgumentException e) {
             Assert.assertTrue(e.getMessage().contains("should not be null"));
         }

         stage.set(stage.paramWithNullDefault, 3);
         Assert.assertEquals(3, (int) stage.get(stage.paramWithNullDefault));
     }

     private static <T> void assertInvalidValue(Stage<?> stage, Param<T> param, T value) {
         try {
             stage.set(param, value);
             Assert.fail("Expected IllegalArgumentException");
         } catch (IllegalArgumentException e) {
             Assert.assertTrue(e.getMessage().contains("invalid value"));
         }
     }

     private static <T> void assertInvalidClass(Stage<?> stage, Param<T> param, Object value) {
         try {
             stage.set(param, (T) value);
             Assert.fail("Expected ClassCastException");
         } catch (ClassCastException e) {
             Assert.assertTrue(e.getMessage().contains("incompatible class"));
         }
     }

     @Test
     public void testSetUndefinedParam() {
         MyStage stage = new MyStage();
         Param<Integer> param = new IntParam("anotherIntParam", "Not defined on MyStage", 1);
         try {
             stage.set(param, 2);
             Assert.fail("Expected IllegalArgumentException");
         } catch (IllegalArgumentException e) {
             Assert.assertTrue(e.getMessage().contains(MyStage.class.getName()));
         }
     }

     @Test
     public void testParamSetInvalidValue() {
         MyStage stage = new MyStage();
         assertInvalidValue(stage, MyParams.INT_PARAM, 100);
         assertInvalidValue(stage, MyParams.LONG_PARAM, 100L);
         assertInvalidValue(stage, MyParams.LONG_PARAM2, Integer.MAX_VALUE + 100L);
         assertInvalidValue(stage, MyParams.FLOAT_PARAM, 100.0f);
         assertInvalidValue(stage, MyParams.DOUBLE_PARAM, 100.0);
         assertInvalidClass(stage, MyParams.INT_PARAM, "100");
         assertInvalidClass(stage, MyParams.STRING_PARAM, 100);

         Param<Integer> param = stage.getParam("stringParam");
         assertInvalidClass(stage, param, 50);
     }

     @Test
     public void testParamSetValidValue() {
         MyStage stage = new MyStage();

         stage.set(MyParams.BOOLEAN_PARAM, true);
         Assert.assertEquals(true, stage.get(MyParams.BOOLEAN_PARAM));

         stage.set(MyParams.INT_PARAM, 50);
         Assert.assertEquals(50, (int) stage.get(MyParams.INT_PARAM));

         stage.set(MyParams.LONG_PARAM, 50L);
         Assert.assertEquals(50L, (long) stage.get(MyParams.LONG_PARAM));

         stage.set(MyParams.LONG_PARAM2, Integer.MAX_VALUE + 50L);
         Assert.assertEquals(Integer.MAX_VALUE + 50L, (long) stage.get(MyParams.LONG_PARAM2));

         stage.set(MyParams.FLOAT_PARAM, 50f);
         Assert.assertEquals(50f, stage.get(MyParams.FLOAT_PARAM), 0.0001);

         stage.set(MyParams.DOUBLE_PARAM, 50.0);
         Assert.assertEquals(50, stage.get(MyParams.DOUBLE_PARAM), 0.0001);

         stage.set(MyParams.STRING_PARAM, "50");
         Assert.assertEquals("50", stage.get(MyParams.STRING_PARAM));

         stage.set(MyParams.INT_ARRAY_PARAM, new Integer[] {50, 51});
         Assert.assertArrayEquals(new Integer[] {50, 51}, stage.get(MyParams.INT_ARRAY_PARAM));

         stage.set(MyParams.LONG_ARRAY_PARAM, new Long[] {50L, 51L});
         Assert.assertArrayEquals(new Long[] {50L, 51L}, stage.get(MyParams.LONG_ARRAY_PARAM));

         stage.set(MyParams.FLOAT_ARRAY_PARAM, new Float[] {50.0f, 51.0f});
         Assert.assertArrayEquals(new Float[] {50.0f, 51.0f}, stage.get(MyParams.FLOAT_ARRAY_PARAM));

         stage.set(MyParams.DOUBLE_ARRAY_PARAM, new Double[] {50.0, 51.0});
         Assert.assertArrayEquals(new Double[] {50.0, 51.0}, stage.get(MyParams.DOUBLE_ARRAY_PARAM));

         stage.set(MyParams.STRING_ARRAY_PARAM, new String[] {"50", "51"});
         Assert.assertArrayEquals(new String[] {"50", "51"}, stage.get(MyParams.STRING_ARRAY_PARAM));

         stage.set(MyParams.VECTOR_PARAM, Vectors.dense(1, 5, 3));
         Assert.assertEquals(Vectors.dense(1, 5, 3), stage.get(MyParams.VECTOR_PARAM));

         stage.set(
                 MyParams.DOUBLE_ARRAY_ARRAY_PARAM,
                 new Double[][] {new Double[] {50.0, 51.0}, new Double[] {52.0, 53.0}});
         Assert.assertEquals(2, stage.get(MyParams.DOUBLE_ARRAY_ARRAY_PARAM).length);
         Assert.assertArrayEquals(
                 new Double[] {50.0, 51.0}, stage.get(MyParams.DOUBLE_ARRAY_ARRAY_PARAM)[0]);
         Assert.assertArrayEquals(
                 new Double[] {52.0, 53.0}, stage.get(MyParams.DOUBLE_ARRAY_ARRAY_PARAM)[1]);

         stage.set(
                 MyParams.STRING_ARRAY_ARRAY_PARAM,
                 new String[][] {
                     new String[] {"50", "51"},
                     new String[] {"52", "53"}
                 });
         Assert.assertEquals(2, stage.get(MyParams.STRING_ARRAY_ARRAY_PARAM).length);
         Assert.assertArrayEquals(
                 new String[] {"50", "51"}, stage.get(MyParams.STRING_ARRAY_ARRAY_PARAM)[0]);
         Assert.assertArrayEquals(
                 new String[] {"52", "53"}, stage.get(MyParams.STRING_ARRAY_ARRAY_PARAM)[1]);
     }

     @Test
     public void testStageSaveLoad() throws IOException {
         StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
         StreamTableEnvironment tEnv = StreamTableEnvironment.create(env);
         MyStage stage = new MyStage();

         stage.set(stage.paramWithNullDefault, 1);
         stage.set(MyParams.BOOLEAN_PARAM, true);
         stage.set(MyParams.INT_PARAM, 50);
         stage.set(MyParams.LONG_PARAM, 50L);
         stage.set(MyParams.LONG_PARAM2, Integer.MAX_VALUE + 50L);
         stage.set(MyParams.FLOAT_PARAM, 50f);
         stage.set(MyParams.DOUBLE_PARAM, 50.0);
         stage.set(MyParams.STRING_PARAM, "50");
         stage.set(MyParams.INT_ARRAY_PARAM, new Integer[] {50, 51});
         stage.set(MyParams.LONG_ARRAY_PARAM, new Long[] {50L, 51L});
         stage.set(MyParams.FLOAT_ARRAY_PARAM, new Float[] {50.0f, 51.0f});
         stage.set(MyParams.DOUBLE_ARRAY_PARAM, new Double[] {50.0, 51.0});
         stage.set(MyParams.STRING_ARRAY_PARAM, new String[] {"50", "51"});
         stage.set(MyParams.VECTOR_PARAM, Vectors.dense(2, 3, 4));
         stage.set(
                 MyParams.DOUBLE_ARRAY_ARRAY_PARAM,
                 new Double[][] {new Double[] {50.0, 51.0}, new Double[] {52.0, 53.0}});
         stage.set(
                 MyParams.STRING_ARRAY_ARRAY_PARAM,
                 new String[][] {
                     new String[] {"50", "51"},
                     new String[] {"52", "53"}
                 });

         Stage<?> loadedStage = validateStageSaveLoad(tEnv, stage, Collections.emptyMap());

         Assert.assertEquals(1, (int) loadedStage.get(stage.paramWithNullDefault));
         Assert.assertEquals(true, loadedStage.get(MyParams.BOOLEAN_PARAM));
         Assert.assertEquals(50, (int) loadedStage.get(MyParams.INT_PARAM));
         Assert.assertEquals(50L, (long) loadedStage.get(MyParams.LONG_PARAM));
         Assert.assertEquals(Integer.MAX_VALUE + 50L, (long) loadedStage.get(MyParams.LONG_PARAM2));
         Assert.assertEquals(50f, loadedStage.get(MyParams.FLOAT_PARAM), 0.0001);
         Assert.assertEquals(50, loadedStage.get(MyParams.DOUBLE_PARAM), 0.0001);
         Assert.assertEquals("50", loadedStage.get(MyParams.STRING_PARAM));
         Assert.assertArrayEquals(new Integer[] {50, 51}, loadedStage.get(MyParams.INT_ARRAY_PARAM));
         Assert.assertArrayEquals(new Long[] {50L, 51L}, loadedStage.get(MyParams.LONG_ARRAY_PARAM));
         Assert.assertArrayEquals(
                 new Float[] {50.0f, 51.0f}, loadedStage.get(MyParams.FLOAT_ARRAY_PARAM));
         Assert.assertArrayEquals(
                 new Double[] {50.0, 51.0}, loadedStage.get(MyParams.DOUBLE_ARRAY_PARAM));
         Assert.assertArrayEquals(
                 new String[] {"50", "51"}, loadedStage.get(MyParams.STRING_ARRAY_PARAM));
         Assert.assertEquals(2, stage.get(MyParams.DOUBLE_ARRAY_ARRAY_PARAM).length);
         Assert.assertArrayEquals(
                 new Double[] {50.0, 51.0}, stage.get(MyParams.DOUBLE_ARRAY_ARRAY_PARAM)[0]);
         Assert.assertArrayEquals(
                 new Double[] {52.0, 53.0}, stage.get(MyParams.DOUBLE_ARRAY_ARRAY_PARAM)[1]);
         Assert.assertEquals(2, stage.get(MyParams.STRING_ARRAY_ARRAY_PARAM).length);
         Assert.assertArrayEquals(
                 new String[] {"50", "51"}, stage.get(MyParams.STRING_ARRAY_ARRAY_PARAM)[0]);
         Assert.assertArrayEquals(
                 new String[] {"52", "53"}, stage.get(MyParams.STRING_ARRAY_ARRAY_PARAM)[1]);
         Assert.assertEquals(Vectors.dense(2, 3, 4), loadedStage.get(MyParams.VECTOR_PARAM));
     }

     @Test
     public void testStageSaveLoadWithParamOverrides() throws IOException {
         StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
         StreamTableEnvironment tEnv = StreamTableEnvironment.create(env);
         MyStage stage = new MyStage();
         stage.set(stage.paramWithNullDefault, 1);
         Stage<?> loadedStage =
                 validateStageSaveLoad(
                         tEnv, stage, Collections.singletonMap("paramWithNullDefault", 10));
         Assert.assertEquals(10, (int) loadedStage.get(stage.paramWithNullDefault));
     }

     @Test
     public void testStageLoadWithoutLoadMethod() throws IOException {
         StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
         StreamTableEnvironment tEnv = StreamTableEnvironment.create(env);
         MyStageWithoutLoad stage = new MyStageWithoutLoad();
         try {
             validateStageSaveLoad(tEnv, stage, Collections.emptyMap());
             Assert.fail("Expected RuntimeException");
         } catch (RuntimeException e) {
             Assert.assertTrue(e.getMessage().contains("not implemented"));
         }
     }

     @Test
     public void testValidators() {
         ParamValidator<Integer> gt = ParamValidators.gt(10);
         Assert.assertFalse(gt.validate(null));
         Assert.assertFalse(gt.validate(5));
         Assert.assertFalse(gt.validate(10));
         Assert.assertTrue(gt.validate(15));

         ParamValidator<Integer> gtEq = ParamValidators.gtEq(10);
         Assert.assertFalse(gtEq.validate(null));
         Assert.assertFalse(gtEq.validate(5));
         Assert.assertTrue(gtEq.validate(10));
         Assert.assertTrue(gtEq.validate(15));

         ParamValidator<Integer> lt = ParamValidators.lt(10);
         Assert.assertFalse(lt.validate(null));
         Assert.assertTrue(lt.validate(5));
         Assert.assertFalse(lt.validate(10));
         Assert.assertFalse(lt.validate(15));

         ParamValidator<Integer> ltEq = ParamValidators.ltEq(10);
         Assert.assertFalse(ltEq.validate(null));
         Assert.assertTrue(ltEq.validate(5));
         Assert.assertTrue(ltEq.validate(10));
         Assert.assertFalse(ltEq.validate(15));

         ParamValidator<Integer> inRangeInclusive = ParamValidators.inRange(5, 15);
         Assert.assertFalse(inRangeInclusive.validate(null));
         Assert.assertFalse(inRangeInclusive.validate(0));
         Assert.assertTrue(inRangeInclusive.validate(5));
         Assert.assertTrue(inRangeInclusive.validate(10));
         Assert.assertTrue(inRangeInclusive.validate(15));
         Assert.assertFalse(inRangeInclusive.validate(20));

         ParamValidator<Integer> inRangeExclusive = ParamValidators.inRange(5, 15, false, false);
         Assert.assertFalse(inRangeExclusive.validate(null));
         Assert.assertFalse(inRangeExclusive.validate(0));
         Assert.assertFalse(inRangeExclusive.validate(5));
         Assert.assertTrue(inRangeExclusive.validate(10));
         Assert.assertFalse(inRangeExclusive.validate(15));
         Assert.assertFalse(inRangeExclusive.validate(20));

         ParamValidator<Integer> inArray = ParamValidators.inArray(1, 2, 3);
         Assert.assertFalse(inArray.validate(null));
         Assert.assertTrue(inArray.validate(1));
         Assert.assertFalse(inArray.validate(0));

         ParamValidator<Integer> notNull = ParamValidators.notNull();
         Assert.assertTrue(notNull.validate(5));
         Assert.assertFalse(notNull.validate(null));

         ParamValidator<Object[]> nonEmptyArray = ParamValidators.nonEmptyArray();
         Assert.assertTrue(nonEmptyArray.validate(new String[] {"1"}));
         Assert.assertFalse(nonEmptyArray.validate(null));
         Assert.assertFalse(nonEmptyArray.validate(new String[0]));

         ParamValidator<String[]> isSubArray = ParamValidators.isSubSet("a", "b", "c");
         Assert.assertFalse(isSubArray.validate(null));
         Assert.assertFalse(isSubArray.validate(new String[] {"c", "v"}));
         Assert.assertTrue(isSubArray.validate(new String[] {"a", "b"}));
         Assert.assertFalse(isSubArray.validate(new String[] {"e", "v"}));
     }
 }
	/*
	* 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.
	*/

	package org.apache.flink.ml.api;

	import org.apache.flink.ml.linalg.Vector;
	import org.apache.flink.ml.linalg.Vectors;
	import org.apache.flink.ml.param.BooleanParam;
	import org.apache.flink.ml.param.DoubleArrayArrayParam;
	import org.apache.flink.ml.param.DoubleArrayParam;
	import org.apache.flink.ml.param.DoubleParam;
	import org.apache.flink.ml.param.FloatArrayParam;
	import org.apache.flink.ml.param.FloatParam;
	import org.apache.flink.ml.param.IntArrayParam;
	import org.apache.flink.ml.param.IntParam;
	import org.apache.flink.ml.param.LongArrayParam;
	import org.apache.flink.ml.param.LongParam;
	import org.apache.flink.ml.param.Param;
	import org.apache.flink.ml.param.ParamValidator;
	import org.apache.flink.ml.param.ParamValidators;
	import org.apache.flink.ml.param.StringArrayArrayParam;
	import org.apache.flink.ml.param.StringArrayParam;
	import org.apache.flink.ml.param.StringParam;
	import org.apache.flink.ml.param.VectorParam;
	import org.apache.flink.ml.param.WithParams;
	import org.apache.flink.ml.util.ParamUtils;
	import org.apache.flink.ml.util.ReadWriteUtils;
	import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
	import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;

	import org.junit.Assert;
	import org.junit.Test;

	import java.io.IOException;
	import java.nio.file.Files;
	import java.util.Collections;
	import java.util.HashMap;
	import java.util.Map;

	/** Tests the behavior of Stage and WithParams. */
	public class StageTest {

	// A WithParams subclass which has one parameter for each pre-defined parameter type.
	private interface MyParams<T> extends WithParams<T> {
	Param<Boolean> BOOLEAN_PARAM = new BooleanParam("booleanParam", "Description", false);

	Param<Integer> INT_PARAM =
	new IntParam("intParam", "Description", 1, ParamValidators.lt(100));

	// This param might be regarded as integer type given a small value.
	Param<Long> LONG_PARAM =
	new LongParam("longParam", "Description", 2L, ParamValidators.lt(100));

	// This param must be regarded as long type given a large enough value.
	Param<Long> LONG_PARAM2 =
	new LongParam(
	"longParam2",
	"Description",
	Integer.MAX_VALUE + 1L,
	ParamValidators.lt(Integer.MAX_VALUE + 100L));

	Param<Float> FLOAT_PARAM =
	new FloatParam("floatParam", "Description", 3.0f, ParamValidators.lt(100));

	Param<Double> DOUBLE_PARAM =
	new DoubleParam("doubleParam", "Description", 4.0, ParamValidators.lt(100));

	Param<String> STRING_PARAM = new StringParam("stringParam", "Description", "5");

	Param<Integer[]> INT_ARRAY_PARAM =
	new IntArrayParam("intArrayParam", "Description", new Integer[] {6, 7});

	Param<Long[]> LONG_ARRAY_PARAM =
	new LongArrayParam(
	"longArrayParam",
	"Description",
	new Long[] {8L, 9L},
	ParamValidators.alwaysTrue());

	Param<Float[]> FLOAT_ARRAY_PARAM =
	new FloatArrayParam("floatArrayParam", "Description", new Float[] {10.0f, 11.0f});

	Param<Double[]> DOUBLE_ARRAY_PARAM =
	new DoubleArrayParam(
	"doubleArrayParam",
	"Description",
	new Double[] {12.0, 13.0},
	ParamValidators.alwaysTrue());

	Param<String[]> STRING_ARRAY_PARAM =
	new StringArrayParam("stringArrayParam", "Description", new String[] {"14", "15"});

	Param<String[][]> STRING_ARRAY_ARRAY_PARAM =
	new StringArrayArrayParam(
	"stringArrayArrayParam",
	"Description",
	new String[][] {new String[] {"14", "15"}});

	Param<Double[][]> DOUBLE_ARRAY_ARRAY_PARAM =
	new DoubleArrayArrayParam(
	"doubleArrayArrayParam",
	"Description",
	new Double[][] {new Double[] {14.0, 15.0}, new Double[] {16.0, 17.0}});

	Param<Vector> VECTOR_PARAM =
	new VectorParam("vectorParam", "Description", Vectors.dense(1.0, 2.0, 3.0));
	}

	/**
	* A Stage subclass which inherits all parameters from MyParams and defines an extra parameter.
	*/
	public static class MyStage implements Stage<MyStage>, MyParams<MyStage> {
	private final Map<Param<?>, Object> paramMap = new HashMap<>();

	public final Param<Integer> extraIntParam =
	new IntParam("extraIntParam", "Description", 20, ParamValidators.alwaysTrue());

	public final Param<Integer> paramWithNullDefault =
	new IntParam(
	"paramWithNullDefault",
	"Must be explicitly set with a non-null value",
	null,
	ParamValidators.notNull());

	public MyStage() {
	ParamUtils.initializeMapWithDefaultValues(paramMap, this);
	}

	@Override
	public Map<Param<?>, Object> getParamMap() {
	return paramMap;
	}

	@Override
	public void save(String path) throws IOException {
	ReadWriteUtils.saveMetadata(this, path);
	}

	public static MyStage load(StreamTableEnvironment tEnv, String path) throws IOException {
	return ReadWriteUtils.loadStageParam(path);
	}
	}

	/** A Stage subclass without the static load() method. */
	public static class MyStageWithoutLoad implements Stage<MyStage>, MyParams<MyStage> {
	private final Map<Param<?>, Object> paramMap = new HashMap<>();

	public MyStageWithoutLoad() {
	ParamUtils.initializeMapWithDefaultValues(paramMap, this);
	}

	@Override
	public void save(String path) throws IOException {
	ReadWriteUtils.saveMetadata(this, path);
	}

	@Override
	public Map<Param<?>, Object> getParamMap() {
	return paramMap;
	}
	}

	// Asserts that m1 and m2 are equivalent.
	private static void assertParamMapEquals(Map<Param<?>, Object> m1, Map<Param<?>, Object> m2) {
	Assert.assertTrue(m1 != null && m2 != null);
	Assert.assertEquals(m1.size(), m2.size());

	for (Map.Entry<Param<?>, Object> entry : m1.entrySet()) {
	Assert.assertTrue(m2.containsKey(entry.getKey()));
	Object v1 = entry.getValue();
	Object v2 = m2.get(entry.getKey());
	if (v1 == null \|\| v2 == null) {
	Assert.assertTrue(v1 == null && v2 == null);
	} else if (v1.getClass().isArray() && v2.getClass().isArray()) {
	Assert.assertArrayEquals((Object[]) v1, (Object[]) v2);
	} else {
	Assert.assertEquals(v1, v2);
	}
	}
	}

	// Saves and loads the given stage. And verifies that the loaded stage has same parameter values
	// as the original stage.
	private static Stage<?> validateStageSaveLoad(
	StreamTableEnvironment tEnv, Stage<?> stage, Map<String, Object> paramOverrides)
	throws IOException {
	for (Map.Entry<String, Object> entry : paramOverrides.entrySet()) {
	Param<?> param = stage.getParam(entry.getKey());
	ReadWriteUtils.setParam(stage, param, entry.getValue());
	}

	String path = Files.createTempDirectory("").toString();
	stage.save(path);
	try {
	stage.save(path);
	Assert.fail("Expected IOException");
	} catch (IOException e) {
	// This is expected.
	}

	Stage<?> loadedStage = ReadWriteUtils.loadStage(tEnv, path);
	for (Map.Entry<String, Object> entry : paramOverrides.entrySet()) {
	Param<?> param = loadedStage.getParam(entry.getKey());
	Assert.assertEquals(entry.getValue(), loadedStage.get(param));
	}
	assertParamMapEquals(stage.getParamMap(), loadedStage.getParamMap());
	return loadedStage;
	}

	@Test
	public void testParamSetValueWithName() {
	MyStage stage = new MyStage();

	Param<Integer> paramA = MyParams.INT_PARAM;
	stage.set(paramA, 2);
	Assert.assertEquals(2, (int) stage.get(paramA));

	Param<Integer> paramB = stage.getParam("intParam");
	stage.set(paramB, 3);
	Assert.assertEquals(3, (int) stage.get(paramB));

	Param<Integer> paramC = stage.getParam("extraIntParam");
	stage.set(paramC, 50);
	Assert.assertEquals(50, (int) stage.get(paramC));
	}

	@Test
	public void testParamWithNullDefault() {
	MyStage stage = new MyStage();
	try {
	stage.get(stage.paramWithNullDefault);
	Assert.fail("Expected IllegalArgumentException");
	} catch (IllegalArgumentException e) {
	Assert.assertTrue(e.getMessage().contains("should not be null"));
	}

	stage.set(stage.paramWithNullDefault, 3);
	Assert.assertEquals(3, (int) stage.get(stage.paramWithNullDefault));
	}

	private static <T> void assertInvalidValue(Stage<?> stage, Param<T> param, T value) {
	try {
	stage.set(param, value);
	Assert.fail("Expected IllegalArgumentException");
	} catch (IllegalArgumentException e) {
	Assert.assertTrue(e.getMessage().contains("invalid value"));
	}
	}

	private static <T> void assertInvalidClass(Stage<?> stage, Param<T> param, Object value) {
	try {
	stage.set(param, (T) value);
	Assert.fail("Expected ClassCastException");
	} catch (ClassCastException e) {
	Assert.assertTrue(e.getMessage().contains("incompatible class"));
	}
	}

	@Test
	public void testSetUndefinedParam() {
	MyStage stage = new MyStage();
	Param<Integer> param = new IntParam("anotherIntParam", "Not defined on MyStage", 1);
	try {
	stage.set(param, 2);
	Assert.fail("Expected IllegalArgumentException");
	} catch (IllegalArgumentException e) {
	Assert.assertTrue(e.getMessage().contains(MyStage.class.getName()));
	}
	}

	@Test
	public void testParamSetInvalidValue() {
	MyStage stage = new MyStage();
	assertInvalidValue(stage, MyParams.INT_PARAM, 100);
	assertInvalidValue(stage, MyParams.LONG_PARAM, 100L);
	assertInvalidValue(stage, MyParams.LONG_PARAM2, Integer.MAX_VALUE + 100L);
	assertInvalidValue(stage, MyParams.FLOAT_PARAM, 100.0f);
	assertInvalidValue(stage, MyParams.DOUBLE_PARAM, 100.0);
	assertInvalidClass(stage, MyParams.INT_PARAM, "100");
	assertInvalidClass(stage, MyParams.STRING_PARAM, 100);

	Param<Integer> param = stage.getParam("stringParam");
	assertInvalidClass(stage, param, 50);
	}

	@Test
	public void testParamSetValidValue() {
	MyStage stage = new MyStage();

	stage.set(MyParams.BOOLEAN_PARAM, true);
	Assert.assertEquals(true, stage.get(MyParams.BOOLEAN_PARAM));

	stage.set(MyParams.INT_PARAM, 50);
	Assert.assertEquals(50, (int) stage.get(MyParams.INT_PARAM));

	stage.set(MyParams.LONG_PARAM, 50L);
	Assert.assertEquals(50L, (long) stage.get(MyParams.LONG_PARAM));

	stage.set(MyParams.LONG_PARAM2, Integer.MAX_VALUE + 50L);
	Assert.assertEquals(Integer.MAX_VALUE + 50L, (long) stage.get(MyParams.LONG_PARAM2));

	stage.set(MyParams.FLOAT_PARAM, 50f);
	Assert.assertEquals(50f, stage.get(MyParams.FLOAT_PARAM), 0.0001);

	stage.set(MyParams.DOUBLE_PARAM, 50.0);
	Assert.assertEquals(50, stage.get(MyParams.DOUBLE_PARAM), 0.0001);

	stage.set(MyParams.STRING_PARAM, "50");
	Assert.assertEquals("50", stage.get(MyParams.STRING_PARAM));

	stage.set(MyParams.INT_ARRAY_PARAM, new Integer[] {50, 51});
	Assert.assertArrayEquals(new Integer[] {50, 51}, stage.get(MyParams.INT_ARRAY_PARAM));

	stage.set(MyParams.LONG_ARRAY_PARAM, new Long[] {50L, 51L});
	Assert.assertArrayEquals(new Long[] {50L, 51L}, stage.get(MyParams.LONG_ARRAY_PARAM));

	stage.set(MyParams.FLOAT_ARRAY_PARAM, new Float[] {50.0f, 51.0f});
	Assert.assertArrayEquals(new Float[] {50.0f, 51.0f}, stage.get(MyParams.FLOAT_ARRAY_PARAM));

	stage.set(MyParams.DOUBLE_ARRAY_PARAM, new Double[] {50.0, 51.0});
	Assert.assertArrayEquals(new Double[] {50.0, 51.0}, stage.get(MyParams.DOUBLE_ARRAY_PARAM));

	stage.set(MyParams.STRING_ARRAY_PARAM, new String[] {"50", "51"});
	Assert.assertArrayEquals(new String[] {"50", "51"}, stage.get(MyParams.STRING_ARRAY_PARAM));

	stage.set(MyParams.VECTOR_PARAM, Vectors.dense(1, 5, 3));
	Assert.assertEquals(Vectors.dense(1, 5, 3), stage.get(MyParams.VECTOR_PARAM));

	stage.set(
	MyParams.DOUBLE_ARRAY_ARRAY_PARAM,
	new Double[][] {new Double[] {50.0, 51.0}, new Double[] {52.0, 53.0}});
	Assert.assertEquals(2, stage.get(MyParams.DOUBLE_ARRAY_ARRAY_PARAM).length);
	Assert.assertArrayEquals(
	new Double[] {50.0, 51.0}, stage.get(MyParams.DOUBLE_ARRAY_ARRAY_PARAM)[0]);
	Assert.assertArrayEquals(
	new Double[] {52.0, 53.0}, stage.get(MyParams.DOUBLE_ARRAY_ARRAY_PARAM)[1]);

	stage.set(
	MyParams.STRING_ARRAY_ARRAY_PARAM,
	new String[][] {
	new String[] {"50", "51"},
	new String[] {"52", "53"}
	});
	Assert.assertEquals(2, stage.get(MyParams.STRING_ARRAY_ARRAY_PARAM).length);
	Assert.assertArrayEquals(
	new String[] {"50", "51"}, stage.get(MyParams.STRING_ARRAY_ARRAY_PARAM)[0]);
	Assert.assertArrayEquals(
	new String[] {"52", "53"}, stage.get(MyParams.STRING_ARRAY_ARRAY_PARAM)[1]);
	}

	@Test
	public void testStageSaveLoad() throws IOException {
	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
	StreamTableEnvironment tEnv = StreamTableEnvironment.create(env);
	MyStage stage = new MyStage();

	stage.set(stage.paramWithNullDefault, 1);
	stage.set(MyParams.BOOLEAN_PARAM, true);
	stage.set(MyParams.INT_PARAM, 50);
	stage.set(MyParams.LONG_PARAM, 50L);
	stage.set(MyParams.LONG_PARAM2, Integer.MAX_VALUE + 50L);
	stage.set(MyParams.FLOAT_PARAM, 50f);
	stage.set(MyParams.DOUBLE_PARAM, 50.0);
	stage.set(MyParams.STRING_PARAM, "50");
	stage.set(MyParams.INT_ARRAY_PARAM, new Integer[] {50, 51});
	stage.set(MyParams.LONG_ARRAY_PARAM, new Long[] {50L, 51L});
	stage.set(MyParams.FLOAT_ARRAY_PARAM, new Float[] {50.0f, 51.0f});
	stage.set(MyParams.DOUBLE_ARRAY_PARAM, new Double[] {50.0, 51.0});
	stage.set(MyParams.STRING_ARRAY_PARAM, new String[] {"50", "51"});
	stage.set(MyParams.VECTOR_PARAM, Vectors.dense(2, 3, 4));
	stage.set(
	MyParams.DOUBLE_ARRAY_ARRAY_PARAM,
	new Double[][] {new Double[] {50.0, 51.0}, new Double[] {52.0, 53.0}});
	stage.set(
	MyParams.STRING_ARRAY_ARRAY_PARAM,
	new String[][] {
	new String[] {"50", "51"},
	new String[] {"52", "53"}
	});

	Stage<?> loadedStage = validateStageSaveLoad(tEnv, stage, Collections.emptyMap());

	Assert.assertEquals(1, (int) loadedStage.get(stage.paramWithNullDefault));
	Assert.assertEquals(true, loadedStage.get(MyParams.BOOLEAN_PARAM));
	Assert.assertEquals(50, (int) loadedStage.get(MyParams.INT_PARAM));
	Assert.assertEquals(50L, (long) loadedStage.get(MyParams.LONG_PARAM));
	Assert.assertEquals(Integer.MAX_VALUE + 50L, (long) loadedStage.get(MyParams.LONG_PARAM2));
	Assert.assertEquals(50f, loadedStage.get(MyParams.FLOAT_PARAM), 0.0001);
	Assert.assertEquals(50, loadedStage.get(MyParams.DOUBLE_PARAM), 0.0001);
	Assert.assertEquals("50", loadedStage.get(MyParams.STRING_PARAM));
	Assert.assertArrayEquals(new Integer[] {50, 51}, loadedStage.get(MyParams.INT_ARRAY_PARAM));
	Assert.assertArrayEquals(new Long[] {50L, 51L}, loadedStage.get(MyParams.LONG_ARRAY_PARAM));
	Assert.assertArrayEquals(
	new Float[] {50.0f, 51.0f}, loadedStage.get(MyParams.FLOAT_ARRAY_PARAM));
	Assert.assertArrayEquals(
	new Double[] {50.0, 51.0}, loadedStage.get(MyParams.DOUBLE_ARRAY_PARAM));
	Assert.assertArrayEquals(
	new String[] {"50", "51"}, loadedStage.get(MyParams.STRING_ARRAY_PARAM));
	Assert.assertEquals(2, stage.get(MyParams.DOUBLE_ARRAY_ARRAY_PARAM).length);
	Assert.assertArrayEquals(
	new Double[] {50.0, 51.0}, stage.get(MyParams.DOUBLE_ARRAY_ARRAY_PARAM)[0]);
	Assert.assertArrayEquals(
	new Double[] {52.0, 53.0}, stage.get(MyParams.DOUBLE_ARRAY_ARRAY_PARAM)[1]);
	Assert.assertEquals(2, stage.get(MyParams.STRING_ARRAY_ARRAY_PARAM).length);
	Assert.assertArrayEquals(
	new String[] {"50", "51"}, stage.get(MyParams.STRING_ARRAY_ARRAY_PARAM)[0]);
	Assert.assertArrayEquals(
	new String[] {"52", "53"}, stage.get(MyParams.STRING_ARRAY_ARRAY_PARAM)[1]);
	Assert.assertEquals(Vectors.dense(2, 3, 4), loadedStage.get(MyParams.VECTOR_PARAM));
	}

	@Test
	public void testStageSaveLoadWithParamOverrides() throws IOException {
	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
	StreamTableEnvironment tEnv = StreamTableEnvironment.create(env);
	MyStage stage = new MyStage();
	stage.set(stage.paramWithNullDefault, 1);
	Stage<?> loadedStage =
	validateStageSaveLoad(
	tEnv, stage, Collections.singletonMap("paramWithNullDefault", 10));
	Assert.assertEquals(10, (int) loadedStage.get(stage.paramWithNullDefault));
	}

	@Test
	public void testStageLoadWithoutLoadMethod() throws IOException {
	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
	StreamTableEnvironment tEnv = StreamTableEnvironment.create(env);
	MyStageWithoutLoad stage = new MyStageWithoutLoad();
	try {
	validateStageSaveLoad(tEnv, stage, Collections.emptyMap());
	Assert.fail("Expected RuntimeException");
	} catch (RuntimeException e) {
	Assert.assertTrue(e.getMessage().contains("not implemented"));
	}
	}

	@Test
	public void testValidators() {
	ParamValidator<Integer> gt = ParamValidators.gt(10);
	Assert.assertFalse(gt.validate(null));
	Assert.assertFalse(gt.validate(5));
	Assert.assertFalse(gt.validate(10));
	Assert.assertTrue(gt.validate(15));

	ParamValidator<Integer> gtEq = ParamValidators.gtEq(10);
	Assert.assertFalse(gtEq.validate(null));
	Assert.assertFalse(gtEq.validate(5));
	Assert.assertTrue(gtEq.validate(10));
	Assert.assertTrue(gtEq.validate(15));

	ParamValidator<Integer> lt = ParamValidators.lt(10);
	Assert.assertFalse(lt.validate(null));
	Assert.assertTrue(lt.validate(5));
	Assert.assertFalse(lt.validate(10));
	Assert.assertFalse(lt.validate(15));

	ParamValidator<Integer> ltEq = ParamValidators.ltEq(10);
	Assert.assertFalse(ltEq.validate(null));
	Assert.assertTrue(ltEq.validate(5));
	Assert.assertTrue(ltEq.validate(10));
	Assert.assertFalse(ltEq.validate(15));

	ParamValidator<Integer> inRangeInclusive = ParamValidators.inRange(5, 15);
	Assert.assertFalse(inRangeInclusive.validate(null));
	Assert.assertFalse(inRangeInclusive.validate(0));
	Assert.assertTrue(inRangeInclusive.validate(5));
	Assert.assertTrue(inRangeInclusive.validate(10));
	Assert.assertTrue(inRangeInclusive.validate(15));
	Assert.assertFalse(inRangeInclusive.validate(20));

	ParamValidator<Integer> inRangeExclusive = ParamValidators.inRange(5, 15, false, false);
	Assert.assertFalse(inRangeExclusive.validate(null));
	Assert.assertFalse(inRangeExclusive.validate(0));
	Assert.assertFalse(inRangeExclusive.validate(5));
	Assert.assertTrue(inRangeExclusive.validate(10));
	Assert.assertFalse(inRangeExclusive.validate(15));
	Assert.assertFalse(inRangeExclusive.validate(20));

	ParamValidator<Integer> inArray = ParamValidators.inArray(1, 2, 3);
	Assert.assertFalse(inArray.validate(null));
	Assert.assertTrue(inArray.validate(1));
	Assert.assertFalse(inArray.validate(0));

	ParamValidator<Integer> notNull = ParamValidators.notNull();
	Assert.assertTrue(notNull.validate(5));
	Assert.assertFalse(notNull.validate(null));

	ParamValidator<Object[]> nonEmptyArray = ParamValidators.nonEmptyArray();
	Assert.assertTrue(nonEmptyArray.validate(new String[] {"1"}));
	Assert.assertFalse(nonEmptyArray.validate(null));
	Assert.assertFalse(nonEmptyArray.validate(new String[0]));

	ParamValidator<String[]> isSubArray = ParamValidators.isSubSet("a", "b", "c");
	Assert.assertFalse(isSubArray.validate(null));
	Assert.assertFalse(isSubArray.validate(new String[] {"c", "v"}));
	Assert.assertTrue(isSubArray.validate(new String[] {"a", "b"}));
	Assert.assertFalse(isSubArray.validate(new String[] {"e", "v"}));
	}
	}