You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by li...@apache.org on 2022/04/14 11:55:46 UTC
[flink-ml] 01/01: [FLINK-27072] Add Transformer for Bucketizer
This is an automated email from the ASF dual-hosted git repository.
lindong pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink-ml.git
commit d8aaa72db26142dfbf6656116bc96e656771630a
Author: zhangzp <zh...@gmail.com>
AuthorDate: Fri Apr 8 19:15:48 2022 +0800
[FLINK-27072] Add Transformer for Bucketizer
---
.../{ArrayParam.java => ArrayArrayParam.java} | 15 +-
.../java/org/apache/flink/ml/param/ArrayParam.java | 2 +-
...{ArrayParam.java => DoubleArrayArrayParam.java} | 24 +--
.../java/org/apache/flink/ml/api/StageTest.java | 24 +++
.../flink/ml/common/param/HasHandleInvalid.java | 3 +-
.../flink/ml/feature/bucketizer/Bucketizer.java | 164 ++++++++++++++++++
.../ml/feature/bucketizer/BucketizerParams.java | 79 +++++++++
.../ml/feature/onehotencoder/OneHotEncoder.java | 4 +
.../feature/onehotencoder/OneHotEncoderModel.java | 4 +
.../feature/onehotencoder/OneHotEncoderParams.java | 3 -
.../ml/feature/stringindexer/StringIndexer.java | 4 +
.../feature/stringindexer/StringIndexerModel.java | 4 +
.../stringindexer/StringIndexerModelParams.java | 3 -
.../feature/vectorassembler/VectorAssembler.java | 4 +
.../vectorassembler/VectorAssemblerParams.java | 3 -
.../ml/classification/LogisticRegressionTest.java | 5 +-
.../flink/ml/classification/NaiveBayesTest.java | 16 +-
.../flink/ml/clustering/OnlineKMeansTest.java | 6 +-
.../apache/flink/ml/feature/BucketizerTest.java | 190 +++++++++++++++++++++
.../apache/flink/ml/feature/OneHotEncoderTest.java | 6 +-
.../flink/ml/feature/StandardScalerTest.java | 6 +-
.../flink/ml/feature/VectorAssemblerTest.java | 8 +-
.../stringindexer/IndexToStringModelTest.java | 8 +-
.../feature/stringindexer/StringIndexerTest.java | 6 +-
24 files changed, 516 insertions(+), 75 deletions(-)
diff --git a/flink-ml-core/src/main/java/org/apache/flink/ml/param/ArrayParam.java b/flink-ml-core/src/main/java/org/apache/flink/ml/param/ArrayArrayParam.java
similarity index 78%
copy from flink-ml-core/src/main/java/org/apache/flink/ml/param/ArrayParam.java
copy to flink-ml-core/src/main/java/org/apache/flink/ml/param/ArrayArrayParam.java
index adc3b80..6b6950d 100644
--- a/flink-ml-core/src/main/java/org/apache/flink/ml/param/ArrayParam.java
+++ b/flink-ml-core/src/main/java/org/apache/flink/ml/param/ArrayArrayParam.java
@@ -22,21 +22,20 @@ import org.apache.flink.ml.util.ReadWriteUtils;
import java.io.IOException;
-/** Class for array-type parameters. */
-public class ArrayParam<T> extends Param<T[]> {
+/** Class for the array parameter. */
+public class ArrayArrayParam<T> extends Param<T[][]> {
- public ArrayParam(
+ public ArrayArrayParam(
String name,
- Class<T[]> clazz,
+ Class<T[][]> clazz,
String description,
- T[] defaultValue,
- ParamValidator<T[]> validator) {
+ T[][] defaultValue,
+ ParamValidator<T[][]> validator) {
super(name, clazz, description, defaultValue, validator);
}
@Override
- public T[] jsonDecode(Object json) throws IOException {
- // Converts Object Mapper's default ArrayList<T> to T[].
+ public T[][] jsonDecode(Object json) throws IOException {
String jsonStr = ReadWriteUtils.OBJECT_MAPPER.writeValueAsString(json);
return ReadWriteUtils.OBJECT_MAPPER.readValue(jsonStr, clazz);
}
diff --git a/flink-ml-core/src/main/java/org/apache/flink/ml/param/ArrayParam.java b/flink-ml-core/src/main/java/org/apache/flink/ml/param/ArrayParam.java
index adc3b80..f859fab 100644
--- a/flink-ml-core/src/main/java/org/apache/flink/ml/param/ArrayParam.java
+++ b/flink-ml-core/src/main/java/org/apache/flink/ml/param/ArrayParam.java
@@ -22,7 +22,7 @@ import org.apache.flink.ml.util.ReadWriteUtils;
import java.io.IOException;
-/** Class for array-type parameters. */
+/** Class for the array parameter. */
public class ArrayParam<T> extends Param<T[]> {
public ArrayParam(
diff --git a/flink-ml-core/src/main/java/org/apache/flink/ml/param/ArrayParam.java b/flink-ml-core/src/main/java/org/apache/flink/ml/param/DoubleArrayArrayParam.java
similarity index 58%
copy from flink-ml-core/src/main/java/org/apache/flink/ml/param/ArrayParam.java
copy to flink-ml-core/src/main/java/org/apache/flink/ml/param/DoubleArrayArrayParam.java
index adc3b80..eef9a6b 100644
--- a/flink-ml-core/src/main/java/org/apache/flink/ml/param/ArrayParam.java
+++ b/flink-ml-core/src/main/java/org/apache/flink/ml/param/DoubleArrayArrayParam.java
@@ -18,26 +18,18 @@
package org.apache.flink.ml.param;
-import org.apache.flink.ml.util.ReadWriteUtils;
+/** Class for the array of double array parameter. */
+public class DoubleArrayArrayParam extends ArrayArrayParam<Double> {
-import java.io.IOException;
-
-/** Class for array-type parameters. */
-public class ArrayParam<T> extends Param<T[]> {
-
- public ArrayParam(
+ public DoubleArrayArrayParam(
String name,
- Class<T[]> clazz,
String description,
- T[] defaultValue,
- ParamValidator<T[]> validator) {
- super(name, clazz, description, defaultValue, validator);
+ Double[][] defaultValue,
+ ParamValidator<Double[][]> validator) {
+ super(name, Double[][].class, description, defaultValue, validator);
}
- @Override
- public T[] jsonDecode(Object json) throws IOException {
- // Converts Object Mapper's default ArrayList<T> to T[].
- String jsonStr = ReadWriteUtils.OBJECT_MAPPER.writeValueAsString(json);
- return ReadWriteUtils.OBJECT_MAPPER.readValue(jsonStr, clazz);
+ public DoubleArrayArrayParam(String name, String description, Double[][] defaultValue) {
+ this(name, description, defaultValue, ParamValidators.alwaysTrue());
}
}
diff --git a/flink-ml-core/src/test/java/org/apache/flink/ml/api/StageTest.java b/flink-ml-core/src/test/java/org/apache/flink/ml/api/StageTest.java
index 091502b..6aa159c 100644
--- a/flink-ml-core/src/test/java/org/apache/flink/ml/api/StageTest.java
+++ b/flink-ml-core/src/test/java/org/apache/flink/ml/api/StageTest.java
@@ -19,6 +19,7 @@
package org.apache.flink.ml.api;
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;
@@ -99,6 +100,12 @@ public class StageTest {
Param<String[]> STRING_ARRAY_PARAM =
new StringArrayParam("stringArrayParam", "Description", 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}});
}
/**
@@ -317,6 +324,15 @@ public class StageTest {
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.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]);
}
@Test
@@ -338,6 +354,9 @@ public class StageTest {
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.DOUBLE_ARRAY_ARRAY_PARAM,
+ new Double[][] {new Double[] {50.0, 51.0}, new Double[] {52.0, 53.0}});
Stage<?> loadedStage = validateStageSaveLoad(tEnv, stage, Collections.emptyMap());
@@ -357,6 +376,11 @@ public class StageTest {
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]);
}
@Test
diff --git a/flink-ml-lib/src/main/java/org/apache/flink/ml/common/param/HasHandleInvalid.java b/flink-ml-lib/src/main/java/org/apache/flink/ml/common/param/HasHandleInvalid.java
index e86ca62..25f15dd 100644
--- a/flink-ml-lib/src/main/java/org/apache/flink/ml/common/param/HasHandleInvalid.java
+++ b/flink-ml-lib/src/main/java/org/apache/flink/ml/common/param/HasHandleInvalid.java
@@ -32,8 +32,7 @@ import org.apache.flink.ml.param.WithParams;
* <ul>
* <li>error: raise an exception.
* <li>skip: filter out rows with bad values.
- * <li>keep: keep bad rows according to a specific rule. Check out each algorithm for detail
- * rules.
+ * <li>keep: keep bad rows according to a specific rule. Check out each algorithm for details.
* </ul>
*/
public interface HasHandleInvalid<T> extends WithParams<T> {
diff --git a/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/bucketizer/Bucketizer.java b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/bucketizer/Bucketizer.java
new file mode 100644
index 0000000..836e30b
--- /dev/null
+++ b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/bucketizer/Bucketizer.java
@@ -0,0 +1,164 @@
+/*
+ * 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.feature.bucketizer;
+
+import org.apache.flink.api.common.functions.FlatMapFunction;
+import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.ml.api.Transformer;
+import org.apache.flink.ml.common.datastream.TableUtils;
+import org.apache.flink.ml.common.param.HasHandleInvalid;
+import org.apache.flink.ml.param.Param;
+import org.apache.flink.ml.util.ParamUtils;
+import org.apache.flink.ml.util.ReadWriteUtils;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.table.api.Table;
+import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
+import org.apache.flink.table.api.internal.TableImpl;
+import org.apache.flink.types.Row;
+import org.apache.flink.util.Collector;
+import org.apache.flink.util.Preconditions;
+
+import org.apache.commons.lang3.ArrayUtils;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Bucketizer is a transformer that maps multiple columns of continuous features to multiple columns
+ * of discrete features, i.e., buckets indices. The indices are in [0, numSplitsInThisColumn - 1].
+ *
+ * <p>The `keep` option of {@link HasHandleInvalid} means that we put the invalid data in the last
+ * bucket of the splits, whose index is the number of the buckets.
+ */
+public class Bucketizer implements Transformer<Bucketizer>, BucketizerParams<Bucketizer> {
+ private final Map<Param<?>, Object> paramMap = new HashMap<>();
+
+ public Bucketizer() {
+ ParamUtils.initializeMapWithDefaultValues(paramMap, this);
+ }
+
+ @Override
+ public Table[] transform(Table... inputs) {
+ Preconditions.checkArgument(inputs.length == 1);
+ String[] inputCols = getInputCols();
+ String[] outputCols = getOutputCols();
+ Double[][] splitsArray = getSplitsArray();
+ Preconditions.checkArgument(inputCols.length == outputCols.length);
+ Preconditions.checkArgument(inputCols.length == splitsArray.length);
+
+ StreamTableEnvironment tEnv =
+ (StreamTableEnvironment) ((TableImpl) inputs[0]).getTableEnvironment();
+
+ RowTypeInfo inputTypeInfo = TableUtils.getRowTypeInfo(inputs[0].getResolvedSchema());
+ TypeInformation<?>[] outputTypes = new TypeInformation[outputCols.length];
+ Arrays.fill(outputTypes, BasicTypeInfo.INT_TYPE_INFO);
+ RowTypeInfo outputTypeInfo =
+ new RowTypeInfo(
+ ArrayUtils.addAll(inputTypeInfo.getFieldTypes(), outputTypes),
+ ArrayUtils.addAll(inputTypeInfo.getFieldNames(), getOutputCols()));
+
+ DataStream<Row> result =
+ tEnv.toDataStream(inputs[0])
+ .flatMap(
+ new FindBucketFunction(inputCols, splitsArray, getHandleInvalid()),
+ outputTypeInfo);
+ return new Table[] {tEnv.fromDataStream(result)};
+ }
+
+ /** Finds the bucket index for each continuous feature of an input data point. */
+ private static class FindBucketFunction implements FlatMapFunction<Row, Row> {
+ private final String[] inputCols;
+ private final String handleInvalid;
+ private final Double[][] splitsArray;
+
+ public FindBucketFunction(
+ String[] inputCols, Double[][] splitsArray, String handleInvalid) {
+ this.inputCols = inputCols;
+ this.splitsArray = splitsArray;
+ this.handleInvalid = handleInvalid;
+ }
+
+ @Override
+ public void flatMap(Row value, Collector<Row> out) {
+ Row outputRow = new Row(inputCols.length);
+
+ for (int i = 0; i < inputCols.length; i++) {
+ double feature = ((Number) value.getField(inputCols[i])).doubleValue();
+ Double[] splits = splitsArray[i];
+ boolean isInvalid = false;
+
+ if (!Double.isNaN(feature)) {
+ int index = Arrays.binarySearch(splits, feature);
+ if (index >= 0) {
+ if (index == inputCols.length - 1) {
+ index--;
+ }
+ outputRow.setField(i, index);
+ } else {
+ index = -index - 1;
+ if (index == 0 || index == inputCols.length) {
+ isInvalid = true;
+ } else {
+ outputRow.setField(i, index - 1);
+ }
+ }
+ } else {
+ isInvalid = true;
+ }
+
+ if (isInvalid) {
+ switch (handleInvalid) {
+ case ERROR_INVALID:
+ throw new RuntimeException(
+ "The input contains invalid value. See "
+ + HANDLE_INVALID
+ + " parameter for more options.");
+ case SKIP_INVALID:
+ return;
+ case KEEP_INVALID:
+ outputRow.setField(i, splits.length - 1);
+ break;
+ default:
+ throw new UnsupportedOperationException(
+ "Unsupported " + HANDLE_INVALID + " type: " + handleInvalid);
+ }
+ }
+ }
+ out.collect(Row.join(value, outputRow));
+ }
+ }
+
+ @Override
+ public void save(String path) throws IOException {
+ ReadWriteUtils.saveMetadata(this, path);
+ }
+
+ public static Bucketizer load(StreamTableEnvironment tEnv, String path) throws IOException {
+ return ReadWriteUtils.loadStageParam(path);
+ }
+
+ @Override
+ public Map<Param<?>, Object> getParamMap() {
+ return paramMap;
+ }
+}
diff --git a/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/bucketizer/BucketizerParams.java b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/bucketizer/BucketizerParams.java
new file mode 100644
index 0000000..34b3918
--- /dev/null
+++ b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/bucketizer/BucketizerParams.java
@@ -0,0 +1,79 @@
+/*
+ * 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.feature.bucketizer;
+
+import org.apache.flink.ml.common.param.HasHandleInvalid;
+import org.apache.flink.ml.common.param.HasInputCols;
+import org.apache.flink.ml.common.param.HasOutputCols;
+import org.apache.flink.ml.param.DoubleArrayArrayParam;
+import org.apache.flink.ml.param.ParamValidator;
+
+/**
+ * Params for {@link Bucketizer}.
+ *
+ * @param <T> The class type of this instance.
+ */
+public interface BucketizerParams<T>
+ extends HasInputCols<T>, HasOutputCols<T>, HasHandleInvalid<T> {
+ /**
+ * The array of split points for mapping continuous features into buckets for multiple columns.
+ *
+ * <p>Each input column is supposed to be mapped into {numberOfSplitPoints - 1} buckets. A
+ * bucket is defined by two split points. For example, bucket(x,y) contains values in the range
+ * [x,y). An exception is that the last bucket also contains y. The array should contain at
+ * least three split points and be strictly increasing.
+ */
+ DoubleArrayArrayParam SPLITS_ARRAY =
+ new DoubleArrayArrayParam(
+ "splitsArray",
+ "Array of split points for mapping continuous features into buckets.",
+ null,
+ new SplitsArrayValidator());
+
+ default Double[][] getSplitsArray() {
+ return get(SPLITS_ARRAY);
+ }
+
+ default T setSplitsArray(Double[][] value) {
+ set(SPLITS_ARRAY, value);
+ return (T) this;
+ }
+
+ /** Param validator for splitsArray. */
+ class SplitsArrayValidator implements ParamValidator<Double[][]> {
+
+ @Override
+ public boolean validate(Double[][] splitsArray) {
+ if (null == splitsArray || splitsArray.length == 0) {
+ return false;
+ }
+ for (Double[] splits : splitsArray) {
+ if (splits.length < 3) {
+ return false;
+ }
+ for (int j = 1; j < splits.length; j++) {
+ if (splits[j] <= splits[j - 1]) {
+ return false;
+ }
+ }
+ }
+ return true;
+ }
+ }
+}
diff --git a/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/onehotencoder/OneHotEncoder.java b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/onehotencoder/OneHotEncoder.java
index b3baf66..9bc1bb7 100644
--- a/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/onehotencoder/OneHotEncoder.java
+++ b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/onehotencoder/OneHotEncoder.java
@@ -23,6 +23,7 @@ import org.apache.flink.api.common.functions.MapPartitionFunction;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.ml.api.Estimator;
import org.apache.flink.ml.common.datastream.DataStreamUtils;
+import org.apache.flink.ml.common.param.HasHandleInvalid;
import org.apache.flink.ml.param.Param;
import org.apache.flink.ml.util.ParamUtils;
import org.apache.flink.ml.util.ReadWriteUtils;
@@ -44,6 +45,9 @@ import java.util.Map;
* <p>Data of selected input columns should be indexed numbers in order for OneHotEncoder to
* function correctly.
*
+ * <p>The `keep` and `skip` option of {@link HasHandleInvalid} is not supported in {@link
+ * OneHotEncoderParams}.
+ *
* <p>See https://en.wikipedia.org/wiki/One-hot.
*/
public class OneHotEncoder
diff --git a/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/onehotencoder/OneHotEncoderModel.java b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/onehotencoder/OneHotEncoderModel.java
index 9bb085f..f9c3e74 100644
--- a/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/onehotencoder/OneHotEncoderModel.java
+++ b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/onehotencoder/OneHotEncoderModel.java
@@ -25,6 +25,7 @@ import org.apache.flink.api.java.typeutils.RowTypeInfo;
import org.apache.flink.ml.api.Model;
import org.apache.flink.ml.common.broadcast.BroadcastUtils;
import org.apache.flink.ml.common.datastream.TableUtils;
+import org.apache.flink.ml.common.param.HasHandleInvalid;
import org.apache.flink.ml.linalg.Vectors;
import org.apache.flink.ml.param.Param;
import org.apache.flink.ml.util.ParamUtils;
@@ -50,6 +51,9 @@ import java.util.function.Function;
/**
* A Model which encodes data into one-hot format using the model data computed by {@link
* OneHotEncoder}.
+ *
+ * <p>The `keep` and `skip` option of {@link HasHandleInvalid} is not supported in {@link
+ * OneHotEncoderParams}.
*/
public class OneHotEncoderModel
implements Model<OneHotEncoderModel>, OneHotEncoderParams<OneHotEncoderModel> {
diff --git a/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/onehotencoder/OneHotEncoderParams.java b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/onehotencoder/OneHotEncoderParams.java
index 4f2347f..9b57159 100644
--- a/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/onehotencoder/OneHotEncoderParams.java
+++ b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/onehotencoder/OneHotEncoderParams.java
@@ -27,9 +27,6 @@ import org.apache.flink.ml.param.Param;
/**
* Params of OneHotEncoderModel.
*
- * <p>The `keep` and `skip` option of {@link HasHandleInvalid} is not supported in {@link
- * OneHotEncoderParams}.
- *
* @param <T> The class type of this instance.
*/
public interface OneHotEncoderParams<T>
diff --git a/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/stringindexer/StringIndexer.java b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/stringindexer/StringIndexer.java
index b1a7881..ee560e0 100644
--- a/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/stringindexer/StringIndexer.java
+++ b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/stringindexer/StringIndexer.java
@@ -25,6 +25,7 @@ import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.ml.api.Estimator;
import org.apache.flink.ml.common.datastream.DataStreamUtils;
+import org.apache.flink.ml.common.param.HasHandleInvalid;
import org.apache.flink.ml.param.Param;
import org.apache.flink.ml.util.ParamUtils;
import org.apache.flink.ml.util.ReadWriteUtils;
@@ -54,6 +55,9 @@ import java.util.Map;
* <p>The input columns are cast to string if they are numeric values. By default, the output model
* is arbitrarily ordered. Users can control this by setting {@link
* StringIndexerParams#STRING_ORDER_TYPE}.
+ *
+ * <p>The `keep` option of {@link HasHandleInvalid} means that we put the invalid entries in a
+ * special bucket, whose index is the number of distinct values in this column.
*/
public class StringIndexer
implements Estimator<StringIndexer, StringIndexerModel>,
diff --git a/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/stringindexer/StringIndexerModel.java b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/stringindexer/StringIndexerModel.java
index 1488610..f1257cc 100644
--- a/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/stringindexer/StringIndexerModel.java
+++ b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/stringindexer/StringIndexerModel.java
@@ -25,6 +25,7 @@ import org.apache.flink.api.java.typeutils.RowTypeInfo;
import org.apache.flink.ml.api.Model;
import org.apache.flink.ml.common.broadcast.BroadcastUtils;
import org.apache.flink.ml.common.datastream.TableUtils;
+import org.apache.flink.ml.common.param.HasHandleInvalid;
import org.apache.flink.ml.param.Param;
import org.apache.flink.ml.util.ParamUtils;
import org.apache.flink.ml.util.ReadWriteUtils;
@@ -47,6 +48,9 @@ import java.util.Map;
/**
* A Model which transforms input string/numeric column(s) to integer column(s) using the model data
* computed by {@link StringIndexer}.
+ *
+ * <p>The `keep` option of {@link HasHandleInvalid} means that we put the invalid entries in a
+ * special bucket, whose index is the number of distinct values in this column.
*/
public class StringIndexerModel
implements Model<StringIndexerModel>, StringIndexerModelParams<StringIndexerModel> {
diff --git a/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/stringindexer/StringIndexerModelParams.java b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/stringindexer/StringIndexerModelParams.java
index 955c672..3f0dd10 100644
--- a/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/stringindexer/StringIndexerModelParams.java
+++ b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/stringindexer/StringIndexerModelParams.java
@@ -25,9 +25,6 @@ import org.apache.flink.ml.common.param.HasOutputCols;
/**
* Params of {@link StringIndexerModel}.
*
- * <p>The `keep` option of {@link HasHandleInvalid} means that we put the invalid entries in a
- * special bucket, whose index is the number of distinct values in this column.
- *
* @param <T> The class type of this instance.
*/
public interface StringIndexerModelParams<T>
diff --git a/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/vectorassembler/VectorAssembler.java b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/vectorassembler/VectorAssembler.java
index a5fa48a..7a90f1a 100644
--- a/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/vectorassembler/VectorAssembler.java
+++ b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/vectorassembler/VectorAssembler.java
@@ -23,6 +23,7 @@ import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.typeutils.RowTypeInfo;
import org.apache.flink.ml.api.Transformer;
import org.apache.flink.ml.common.datastream.TableUtils;
+import org.apache.flink.ml.common.param.HasHandleInvalid;
import org.apache.flink.ml.linalg.DenseVector;
import org.apache.flink.ml.linalg.SparseVector;
import org.apache.flink.ml.linalg.Vector;
@@ -47,6 +48,9 @@ import java.util.Map;
/**
* A feature transformer that combines a given list of input columns into a vector column. Types of
* input columns must be either vector or numerical value.
+ *
+ * <p>The `keep` option of {@link HasHandleInvalid} means that we output bad rows with output column
+ * set to null.
*/
public class VectorAssembler
implements Transformer<VectorAssembler>, VectorAssemblerParams<VectorAssembler> {
diff --git a/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/vectorassembler/VectorAssemblerParams.java b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/vectorassembler/VectorAssemblerParams.java
index 2409a1e..aaada93 100644
--- a/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/vectorassembler/VectorAssemblerParams.java
+++ b/flink-ml-lib/src/main/java/org/apache/flink/ml/feature/vectorassembler/VectorAssemblerParams.java
@@ -25,9 +25,6 @@ import org.apache.flink.ml.common.param.HasOutputCol;
/**
* Params of VectorAssembler.
*
- * <p>The `keep` option of {@link HasHandleInvalid} means that we output bad rows with output column
- * set to null.
- *
* @param <T> The class type of this instance.
*/
public interface VectorAssemblerParams<T>
diff --git a/flink-ml-lib/src/test/java/org/apache/flink/ml/classification/LogisticRegressionTest.java b/flink-ml-lib/src/test/java/org/apache/flink/ml/classification/LogisticRegressionTest.java
index a33c9b4..88882c6 100644
--- a/flink-ml-lib/src/test/java/org/apache/flink/ml/classification/LogisticRegressionTest.java
+++ b/flink-ml-lib/src/test/java/org/apache/flink/ml/classification/LogisticRegressionTest.java
@@ -37,6 +37,7 @@ import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
import org.apache.flink.types.Row;
import org.apache.commons.collections.IteratorUtils;
+import org.apache.commons.lang3.exception.ExceptionUtils;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
@@ -275,10 +276,10 @@ public class LogisticRegressionTest {
new LogisticRegression().fit(multinomialDataTable);
env.execute();
fail();
- } catch (Exception e) {
+ } catch (Throwable e) {
assertEquals(
"Multinomial classification is not supported yet. Supported options: [auto, binomial].",
- e.getCause().getCause().getMessage());
+ ExceptionUtils.getRootCause(e).getMessage());
}
}
diff --git a/flink-ml-lib/src/test/java/org/apache/flink/ml/classification/NaiveBayesTest.java b/flink-ml-lib/src/test/java/org/apache/flink/ml/classification/NaiveBayesTest.java
index d19ae14..c481fbb 100644
--- a/flink-ml-lib/src/test/java/org/apache/flink/ml/classification/NaiveBayesTest.java
+++ b/flink-ml-lib/src/test/java/org/apache/flink/ml/classification/NaiveBayesTest.java
@@ -34,6 +34,7 @@ import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
import org.apache.flink.types.Row;
import org.apache.flink.util.CloseableIterator;
+import org.apache.commons.lang3.exception.ExceptionUtils;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Rule;
@@ -198,10 +199,7 @@ public class NaiveBayesTest {
outputTable.execute().collect().next();
Assert.fail("Expected NullPointerException");
} catch (Exception e) {
- Throwable exception = e;
- while (exception.getCause() != null) {
- exception = exception.getCause();
- }
+ Throwable exception = ExceptionUtils.getRootCause(e);
assertEquals(
NaiveBayesModel.class.getName(), exception.getStackTrace()[0].getClassName());
assertEquals("calculateProb", exception.getStackTrace()[0].getMethodName());
@@ -226,10 +224,7 @@ public class NaiveBayesTest {
outputTable.execute().collect().next();
Assert.fail("Expected IllegalArgumentException");
} catch (Exception e) {
- Throwable exception = e;
- while (exception.getCause() != null) {
- exception = exception.getCause();
- }
+ Throwable exception = ExceptionUtils.getRootCause(e);
assertEquals(IllegalArgumentException.class, exception.getClass());
assertEquals("Feature vectors should be of equal length.", exception.getMessage());
}
@@ -249,10 +244,7 @@ public class NaiveBayesTest {
outputTable.execute().collect().next();
Assert.fail("Expected IllegalArgumentException");
} catch (Exception e) {
- Throwable exception = e;
- while (exception.getCause() != null) {
- exception = exception.getCause();
- }
+ Throwable exception = ExceptionUtils.getRootCause(e);
assertEquals(IllegalArgumentException.class, exception.getClass());
assertEquals("Feature vectors should be of equal length.", exception.getMessage());
}
diff --git a/flink-ml-lib/src/test/java/org/apache/flink/ml/clustering/OnlineKMeansTest.java b/flink-ml-lib/src/test/java/org/apache/flink/ml/clustering/OnlineKMeansTest.java
index d13beb0..047e8ea 100644
--- a/flink-ml-lib/src/test/java/org/apache/flink/ml/clustering/OnlineKMeansTest.java
+++ b/flink-ml-lib/src/test/java/org/apache/flink/ml/clustering/OnlineKMeansTest.java
@@ -372,11 +372,7 @@ public class OnlineKMeansTest extends TestLogger {
try {
onlineKMeans.fit(onlineTrainTable);
Assert.fail("Expected IllegalStateException");
- } catch (Exception e) {
- Throwable exception = e;
- while (exception.getCause() != null) {
- exception = exception.getCause();
- }
+ } catch (Throwable exception) {
assertEquals(IllegalStateException.class, exception.getClass());
assertEquals(
"There are more subtasks in the training process than the number "
diff --git a/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/BucketizerTest.java b/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/BucketizerTest.java
new file mode 100644
index 0000000..4a7306c
--- /dev/null
+++ b/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/BucketizerTest.java
@@ -0,0 +1,190 @@
+/*
+ * 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.feature;
+
+import org.apache.flink.api.common.restartstrategy.RestartStrategies;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.ml.common.param.HasHandleInvalid;
+import org.apache.flink.ml.feature.bucketizer.Bucketizer;
+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.Table;
+import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
+import org.apache.flink.test.util.AbstractTestBase;
+import org.apache.flink.types.Row;
+
+import org.apache.commons.collections.IteratorUtils;
+import org.apache.commons.lang3.exception.ExceptionUtils;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.List;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+/** Tests the {@link Bucketizer}. */
+public class BucketizerTest extends AbstractTestBase {
+ @Rule public final TemporaryFolder tempFolder = new TemporaryFolder();
+ private StreamTableEnvironment tEnv;
+ private Table inputTable;
+
+ private static final List<Row> inputData =
+ Arrays.asList(
+ Row.of(1, -0.5, 0.0, 1.0),
+ Row.of(2, Double.NEGATIVE_INFINITY, 1.0, Double.POSITIVE_INFINITY),
+ Row.of(3, Double.NaN, -0.5, -0.5));
+
+ private static final Double[][] splitsArray =
+ new Double[][] {
+ new Double[] {-0.5, 0.0, 0.5},
+ new Double[] {-1.0, 0.0, 2.0},
+ new Double[] {Double.NEGATIVE_INFINITY, 10.0, Double.POSITIVE_INFINITY}
+ };
+
+ private final List<Row> expectedKeepResult =
+ Arrays.asList(Row.of(1, 0, 1, 0), Row.of(2, 2, 1, 1), Row.of(3, 2, 0, 0));
+
+ private final List<Row> expectedSkipResult = Collections.singletonList(Row.of(1, 0, 1, 0));
+
+ @Before
+ public void before() {
+ Configuration config = new Configuration();
+ config.set(ExecutionCheckpointingOptions.ENABLE_CHECKPOINTS_AFTER_TASKS_FINISH, true);
+ StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(config);
+ env.setParallelism(4);
+ env.enableCheckpointing(100);
+ env.setRestartStrategy(RestartStrategies.noRestart());
+ tEnv = StreamTableEnvironment.create(env);
+ inputTable = tEnv.fromDataStream(env.fromCollection(inputData)).as("id", "f1", "f2", "f3");
+ }
+
+ @SuppressWarnings("all")
+ private void verifyOutputResult(Table output, String[] outputCols, List<Row> expectedResult)
+ throws Exception {
+ List<Row> collectedResult =
+ IteratorUtils.toList(tEnv.toDataStream(output).executeAndCollect());
+ List<Row> result = new ArrayList<>(collectedResult.size());
+
+ for (Row r : collectedResult) {
+ Row outRow = new Row(outputCols.length + 1);
+ outRow.setField(0, r.getField("id"));
+ for (int i = 0; i < outputCols.length; i++) {
+ outRow.setField(i + 1, r.getField(outputCols[i]));
+ }
+ result.add(outRow);
+ }
+
+ compareResultCollections(
+ expectedResult, result, Comparator.comparingInt(r -> ((Integer) r.getField(0))));
+ }
+
+ @Test
+ public void testParam() {
+ Bucketizer bucketizer = new Bucketizer();
+ assertEquals(HasHandleInvalid.ERROR_INVALID, bucketizer.getHandleInvalid());
+
+ bucketizer
+ .setInputCols("f1", "f2", "f3")
+ .setOutputCols("o1", "o2", "o3")
+ .setHandleInvalid(HasHandleInvalid.SKIP_INVALID)
+ .setSplitsArray(splitsArray);
+ assertArrayEquals(new String[] {"f1", "f2", "f3"}, bucketizer.getInputCols());
+ assertArrayEquals(new String[] {"o1", "o2", "o3"}, bucketizer.getOutputCols());
+ assertEquals(HasHandleInvalid.SKIP_INVALID, bucketizer.getHandleInvalid());
+
+ Double[][] setSplitsArray = bucketizer.getSplitsArray();
+ assertEquals(splitsArray.length, setSplitsArray.length);
+ for (int i = 0; i < splitsArray.length; i++) {
+ assertArrayEquals(splitsArray[i], setSplitsArray[i]);
+ }
+ }
+
+ @Test
+ public void testOutputSchema() {
+ Bucketizer bucketizer =
+ new Bucketizer()
+ .setInputCols("f1", "f2", "f3")
+ .setOutputCols("o1", "o2", "o3")
+ .setHandleInvalid(HasHandleInvalid.SKIP_INVALID)
+ .setSplitsArray(splitsArray);
+ Table output = bucketizer.transform(inputTable)[0];
+ assertEquals(
+ Arrays.asList("id", "f1", "f2", "f3", "o1", "o2", "o3"),
+ output.getResolvedSchema().getColumnNames());
+ }
+
+ @Test
+ public void testTransform() throws Exception {
+ Bucketizer bucketizer =
+ new Bucketizer()
+ .setInputCols("f1", "f2", "f3")
+ .setOutputCols("o1", "o2", "o3")
+ .setSplitsArray(splitsArray);
+
+ Table output;
+
+ // Tests skip.
+ bucketizer.setHandleInvalid(HasHandleInvalid.SKIP_INVALID);
+ output = bucketizer.transform(inputTable)[0];
+ verifyOutputResult(output, bucketizer.getOutputCols(), expectedSkipResult);
+
+ // Tests keep.
+ bucketizer.setHandleInvalid(HasHandleInvalid.KEEP_INVALID);
+ output = bucketizer.transform(inputTable)[0];
+ verifyOutputResult(output, bucketizer.getOutputCols(), expectedKeepResult);
+
+ // Tests error.
+ bucketizer.setHandleInvalid(HasHandleInvalid.ERROR_INVALID);
+ output = bucketizer.transform(inputTable)[0];
+ try {
+ IteratorUtils.toList(output.execute().collect());
+ fail();
+ } catch (Throwable e) {
+ assertEquals(
+ "The input contains invalid value. See "
+ + HasHandleInvalid.HANDLE_INVALID
+ + " parameter for more options.",
+ ExceptionUtils.getRootCause(e).getMessage());
+ }
+ }
+
+ @Test
+ public void testSaveLoadAndTransform() throws Exception {
+ Bucketizer bucketizer =
+ new Bucketizer()
+ .setInputCols("f1", "f2", "f3")
+ .setOutputCols("o1", "o2", "o3")
+ .setHandleInvalid(HasHandleInvalid.KEEP_INVALID)
+ .setSplitsArray(splitsArray);
+ Bucketizer loadedBucketizer =
+ StageTestUtils.saveAndReload(
+ tEnv, bucketizer, tempFolder.newFolder().getAbsolutePath());
+ Table output = loadedBucketizer.transform(inputTable)[0];
+ verifyOutputResult(output, loadedBucketizer.getOutputCols(), expectedKeepResult);
+ }
+}
diff --git a/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/OneHotEncoderTest.java b/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/OneHotEncoderTest.java
index 85034e0..b7a0758 100644
--- a/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/OneHotEncoderTest.java
+++ b/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/OneHotEncoderTest.java
@@ -36,6 +36,7 @@ import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
import org.apache.flink.types.Row;
import org.apache.flink.util.CloseableIterator;
+import org.apache.commons.lang3.exception.ExceptionUtils;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Rule;
@@ -225,10 +226,7 @@ public class OneHotEncoderTest {
outputTable.execute().collect().next();
Assert.fail("Expected IllegalArgumentException");
} catch (Exception e) {
- Throwable exception = e;
- while (exception.getCause() != null) {
- exception = exception.getCause();
- }
+ Throwable exception = ExceptionUtils.getRootCause(e);
assertEquals(IllegalArgumentException.class, exception.getClass());
assertEquals("Value 0.5 cannot be parsed as indexed integer.", exception.getMessage());
}
diff --git a/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/StandardScalerTest.java b/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/StandardScalerTest.java
index fd0a049..5337973 100644
--- a/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/StandardScalerTest.java
+++ b/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/StandardScalerTest.java
@@ -36,6 +36,7 @@ import org.apache.flink.test.util.AbstractTestBase;
import org.apache.flink.types.Row;
import org.apache.commons.collections.IteratorUtils;
+import org.apache.commons.lang3.exception.ExceptionUtils;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
@@ -267,10 +268,7 @@ public class StandardScalerTest extends AbstractTestBase {
StandardScalerModelData.getModelDataStream(modelDataTable).executeAndCollect());
fail();
} catch (Throwable e) {
- while (e.getCause() != null) {
- e = e.getCause();
- }
- assertEquals("The training set is empty.", e.getMessage());
+ assertEquals("The training set is empty.", ExceptionUtils.getRootCause(e).getMessage());
}
}
}
diff --git a/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/VectorAssemblerTest.java b/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/VectorAssemblerTest.java
index a5f348a..15fd5d9 100644
--- a/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/VectorAssemblerTest.java
+++ b/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/VectorAssemblerTest.java
@@ -35,6 +35,7 @@ import org.apache.flink.test.util.AbstractTestBase;
import org.apache.flink.types.Row;
import org.apache.commons.collections.IteratorUtils;
+import org.apache.commons.lang3.exception.ExceptionUtils;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
@@ -141,10 +142,9 @@ public class VectorAssemblerTest extends AbstractTestBase {
outputTable.execute().collect().next();
Assert.fail("Expected IllegalArgumentException");
} catch (Throwable e) {
- while (e.getCause() != null) {
- e = e.getCause();
- }
- assertEquals("Input column value should not be null.", e.getMessage());
+ assertEquals(
+ "Input column value should not be null.",
+ ExceptionUtils.getRootCause(e).getMessage());
}
}
diff --git a/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/stringindexer/IndexToStringModelTest.java b/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/stringindexer/IndexToStringModelTest.java
index 61c8df4..bc482a1 100644
--- a/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/stringindexer/IndexToStringModelTest.java
+++ b/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/stringindexer/IndexToStringModelTest.java
@@ -29,6 +29,7 @@ import org.apache.flink.test.util.AbstractTestBase;
import org.apache.flink.types.Row;
import org.apache.commons.collections.IteratorUtils;
+import org.apache.commons.lang3.exception.ExceptionUtils;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
@@ -107,10 +108,9 @@ public class IndexToStringModelTest extends AbstractTestBase {
IteratorUtils.toList(tEnv.toDataStream(output).executeAndCollect());
fail();
} catch (Throwable e) {
- while (e.getCause() != null) {
- e = e.getCause();
- }
- assertEquals("The input contains unseen index: 4.", e.getMessage());
+ assertEquals(
+ "The input contains unseen index: 4.",
+ ExceptionUtils.getRootCause(e).getMessage());
}
}
diff --git a/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/stringindexer/StringIndexerTest.java b/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/stringindexer/StringIndexerTest.java
index f011d2a..c8af00e 100644
--- a/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/stringindexer/StringIndexerTest.java
+++ b/flink-ml-lib/src/test/java/org/apache/flink/ml/feature/stringindexer/StringIndexerTest.java
@@ -31,6 +31,7 @@ import org.apache.flink.test.util.AbstractTestBase;
import org.apache.flink.types.Row;
import org.apache.commons.collections.IteratorUtils;
+import org.apache.commons.lang3.exception.ExceptionUtils;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
@@ -218,15 +219,12 @@ public class StringIndexerTest extends AbstractTestBase {
IteratorUtils.toList(tEnv.toDataStream(output).executeAndCollect());
fail();
} catch (Throwable e) {
- while (e.getCause() != null) {
- e = e.getCause();
- }
assertEquals(
"The input contains unseen string: e. "
+ "See "
+ HasHandleInvalid.HANDLE_INVALID
+ " parameter for more options.",
- e.getMessage());
+ ExceptionUtils.getRootCause(e).getMessage());
}
}