yunfengzhou-hub commented on a change in pull request #32: URL: https://github.com/apache/flink-ml/pull/32#discussion_r757201463
########## File path: flink-ml-lib/src/test/java/org/apache/flink/ml/classification/NaiveBayesTest.java ########## @@ -0,0 +1,327 @@ +/* + * 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.classification; + +import org.apache.flink.api.common.eventtime.WatermarkStrategy; +import org.apache.flink.api.common.restartstrategy.RestartStrategies; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.ml.classification.naivebayes.NaiveBayes; +import org.apache.flink.ml.classification.naivebayes.NaiveBayesModel; +import org.apache.flink.ml.classification.naivebayes.NaiveBayesModelData; +import org.apache.flink.ml.linalg.DenseVector; +import org.apache.flink.ml.linalg.Vector; +import org.apache.flink.ml.linalg.Vectors; +import org.apache.flink.ml.util.ReadWriteUtils; +import org.apache.flink.ml.util.StageTestUtils; +import org.apache.flink.streaming.api.environment.ExecutionCheckpointingOptions; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.table.api.DataTypes; +import org.apache.flink.table.api.Schema; +import org.apache.flink.table.api.Table; +import org.apache.flink.table.api.bridge.java.StreamTableEnvironment; +import org.apache.flink.types.Row; +import org.apache.flink.util.CloseableIterator; + +import org.junit.Before; +import org.junit.Test; + +import java.util.HashMap; +import java.util.Map; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; + +/** Tests {@link NaiveBayes} and {@link NaiveBayesModel}. */ +public class NaiveBayesTest { + private StreamExecutionEnvironment env; + private StreamTableEnvironment tEnv; + private Table trainTable; + private Table predictTable; + private Map<Vector, Integer> expectedOutput; + private NaiveBayes estimator; + + @Before + public void before() { + Configuration config = new Configuration(); + config.set(ExecutionCheckpointingOptions.ENABLE_CHECKPOINTS_AFTER_TASKS_FINISH, true); + env = StreamExecutionEnvironment.getExecutionEnvironment(config); + env.setParallelism(4); + env.enableCheckpointing(100); + env.setRestartStrategy(RestartStrategies.noRestart()); + tEnv = StreamTableEnvironment.create(env); + + Schema trainSchema = + Schema.newBuilder() + .column("f0", DataTypes.of(DenseVector.class)) + .column("f1", DataTypes.INT()) + .columnByMetadata("rowtime", "TIMESTAMP_LTZ(3)") + .watermark("rowtime", "SOURCE_WATERMARK()") + .build(); + + trainTable = + tEnv.fromDataStream( + env.fromElements( + Row.of(Vectors.dense(0, 0.), 11), + Row.of(Vectors.dense(1, 0), 10), + Row.of(Vectors.dense(1, 1.), 10)) + .assignTimestampsAndWatermarks( + WatermarkStrategy.noWatermarks()), + trainSchema) + .as("features", "label"); + + Schema predictSchema = + Schema.newBuilder() + .column("f0", DataTypes.of(DenseVector.class)) + .columnByMetadata("rowtime", "TIMESTAMP_LTZ(3)") + .watermark("rowtime", "SOURCE_WATERMARK()") + .build(); + + predictTable = + tEnv.fromDataStream( + env.fromElements( + Row.of(Vectors.dense(0, 1.)), + Row.of(Vectors.dense(0, 0.)), + Row.of(Vectors.dense(1, 0)), + Row.of(Vectors.dense(1, 1.))) + .assignTimestampsAndWatermarks( + WatermarkStrategy.noWatermarks()), + predictSchema) + .as("features"); + + expectedOutput = + new HashMap<Vector, Integer>() { + { + put(Vectors.dense(0, 1.), 11); + put(Vectors.dense(0, 0.), 11); + put(Vectors.dense(1, 0.), 10); + put(Vectors.dense(1, 1.), 10); + } + }; + + estimator = + new NaiveBayes() + .setSmoothing(1.0) + .setFeaturesCol("features") + .setLabelCol("label") + .setPredictionCol("predict") + .setModelType("multinomial"); + } + + /** + * Executes a given table and collect its results. Results are returned as a map whose key is + * the feature, value is the prediction result. + * + * @param table A table to be executed and to have its result collected + * @param featuresCol Name of the column in the table that contains the features + * @param predictionCol Name of the column in the table that contains the prediction result + * @return A map containing the collected results + */ + private static Map<Vector, Integer> executeAndCollect( + Table table, String featuresCol, String predictionCol) { + Map<Vector, Integer> map = new HashMap<>(); + for (CloseableIterator<Row> it = table.execute().collect(); it.hasNext(); ) { + Row row = it.next(); + map.put( + (Vector) row.getField(featuresCol), + ((Number) row.getField(predictionCol)).intValue()); + } + + return map; + } + + @Test + public void testParam() { + NaiveBayes estimator = new NaiveBayes(); + + assertEquals("features", estimator.getFeaturesCol()); + assertEquals("label", estimator.getLabelCol()); + assertEquals("multinomial", estimator.getModelType()); + assertEquals("prediction", estimator.getPredictionCol()); + assertEquals(1.0, estimator.getSmoothing(), 1e-5); + + estimator + .setFeaturesCol("test_feature") + .setLabelCol("test_label") + .setPredictionCol("test_prediction") + .setSmoothing(2.0); + + assertEquals("test_feature", estimator.getFeaturesCol()); + assertEquals("test_label", estimator.getLabelCol()); + assertEquals("test_prediction", estimator.getPredictionCol()); + assertEquals(2.0, estimator.getSmoothing(), 1e-5); + + NaiveBayesModel model = new NaiveBayesModel(); + + assertEquals("features", model.getFeaturesCol()); + assertEquals("multinomial", model.getModelType()); + assertEquals("prediction", model.getPredictionCol()); + + model.setFeaturesCol("test_feature").setPredictionCol("test_prediction"); + + assertEquals("test_feature", model.getFeaturesCol()); + assertEquals("test_prediction", model.getPredictionCol()); + } + + @Test + public void testFitAndPredict() throws Exception { + NaiveBayesModel model = estimator.fit(trainTable); + Table outputTable = model.transform(predictTable)[0]; + Map<Vector, Integer> actualOutput = + executeAndCollect(outputTable, model.getFeaturesCol(), model.getPredictionCol()); + assertEquals(expectedOutput, actualOutput); + } + + @Test(expected = Exception.class) + public void testPredictUnseenFeature() throws Exception { Review comment: I think `testPredictUnseenFeature` should fail naturally while the program cannot find corresponding value of the feature. Adding a check and set detailed error message might have a bit influence on the performance of this algorithm. Other tests that expect an exception do need to check the message and I'll make this change on those tests. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: issues-unsubscr...@flink.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
