Blacklist a number of prediction field names. (#861) (#868)

elastic · Dec 1, 2019 · f670c04 · f670c04
1 parent 7275f29
commit f670c04
Show file tree

Hide file tree

Showing 6 changed files with 82 additions and 0 deletions.
diff --git a/docs/CHANGELOG.asciidoc b/docs/CHANGELOG.asciidoc
@@ -46,6 +46,8 @@ tree which is trained for both regression and classification. (See {ml-pull}811[
 
 === Bug Fixes
 * Fixes potential memory corruption when determining seasonality. (See {ml-pull}852[#852].)
+* Prevent prediction_field_name clashing with other fields in ml results.
+(See {ml-pull}861[#861].)
 
 
 == {es} version 7.5.0

diff --git a/lib/api/CDataFrameTrainBoostedTreeClassifierRunner.cc b/lib/api/CDataFrameTrainBoostedTreeClassifierRunner.cc
@@ -68,6 +68,13 @@ CDataFrameTrainBoostedTreeClassifierRunner::CDataFrameTrainBoostedTreeClassifier
                   this->dependentVariableFieldName()) == categoricalFieldNames.end()) {
         HANDLE_FATAL(<< "Input error: trying to perform classification with numeric target.");
     }
+    const std::set<std::string> predictionFieldNameBlacklist{
+        IS_TRAINING_FIELD_NAME, PREDICTION_PROBABILITY_FIELD_NAME, TOP_CLASSES_FIELD_NAME};
+    if (predictionFieldNameBlacklist.count(this->predictionFieldName()) > 0) {
+        HANDLE_FATAL(<< "Input error: prediction_field_name must not be equal to any of "
+                     << core::CContainerPrinter::print(predictionFieldNameBlacklist)
+                     << ".");
+    }
 }
 
 CDataFrameTrainBoostedTreeClassifierRunner::CDataFrameTrainBoostedTreeClassifierRunner(

diff --git a/lib/api/CDataFrameTrainBoostedTreeRegressionRunner.cc b/lib/api/CDataFrameTrainBoostedTreeRegressionRunner.cc
@@ -52,6 +52,12 @@ CDataFrameTrainBoostedTreeRegressionRunner::CDataFrameTrainBoostedTreeRegression
                   this->dependentVariableFieldName()) != categoricalFieldNames.end()) {
         HANDLE_FATAL(<< "Input error: trying to perform regression with categorical target.");
     }
+    const std::set<std::string> predictionFieldNameBlacklist{IS_TRAINING_FIELD_NAME};
+    if (predictionFieldNameBlacklist.count(this->predictionFieldName()) > 0) {
+        HANDLE_FATAL(<< "Input error: prediction_field_name must not be equal to any of "
+                     << core::CContainerPrinter::print(predictionFieldNameBlacklist)
+                     << ".");
+    }
 }
 
 CDataFrameTrainBoostedTreeRegressionRunner::CDataFrameTrainBoostedTreeRegressionRunner(

diff --git a/lib/api/unittest/CDataFrameTrainBoostedTreeClassifierRunnerTest.cc b/lib/api/unittest/CDataFrameTrainBoostedTreeClassifierRunnerTest.cc
@@ -25,6 +25,26 @@ using TStrVec = std::vector<std::string>;
 using TStrVecVec = std::vector<TStrVec>;
 }
 
+BOOST_AUTO_TEST_CASE(testPredictionFieldNameClash) {
+    TStrVec errors;
+    auto errorHandler = [&errors](std::string error) { errors.push_back(error); };
+    core::CLogger::CScopeSetFatalErrorHandler scope{errorHandler};
+
+    const auto spec{test::CDataFrameAnalysisSpecificationFactory::predictionSpec(
+        "classification", "dep_var", 5, 6, 13000000, 0, 0, {"dep_var"})};
+    rapidjson::Document jsonParameters;
+    jsonParameters.Parse("{"
+                         "  \"dependent_variable\": \"dep_var\","
+                         "  \"prediction_field_name\": \"is_training\""
+                         "}");
+    const auto parameters{
+        api::CDataFrameTrainBoostedTreeClassifierRunner::parameterReader().read(jsonParameters)};
+    api::CDataFrameTrainBoostedTreeClassifierRunner runner(*spec, parameters);
+
+    BOOST_TEST_REQUIRE(errors.size() == 1);
+    BOOST_TEST_REQUIRE(errors[0] == "Input error: prediction_field_name must not be equal to any of [is_training, prediction_probability, top_classes].");
+}
+
 BOOST_AUTO_TEST_CASE(testWriteOneRow) {
     // Prepare input data frame
     const TStrVec columnNames{"x1", "x2", "x3", "x4", "x5", "x5_prediction"};

diff --git a/lib/api/unittest/CDataFrameTrainBoostedTreeRegressionRunnerTest.cc b/lib/api/unittest/CDataFrameTrainBoostedTreeRegressionRunnerTest.cc
@@ -0,0 +1,46 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+#include <core/CDataFrame.h>
+
+#include <api/CDataFrameAnalysisConfigReader.h>
+#include <api/CDataFrameTrainBoostedTreeRegressionRunner.h>
+
+#include <test/CDataFrameAnalysisSpecificationFactory.h>
+
+#include <boost/test/unit_test.hpp>
+
+#include <string>
+#include <vector>
+
+BOOST_AUTO_TEST_SUITE(CDataFrameTrainBoostedTreeRegressionRunnerTest)
+
+using namespace ml;
+namespace {
+using TStrVec = std::vector<std::string>;
+}
+
+BOOST_AUTO_TEST_CASE(testPredictionFieldNameClash) {
+    TStrVec errors;
+    auto errorHandler = [&errors](std::string error) { errors.push_back(error); };
+    core::CLogger::CScopeSetFatalErrorHandler scope{errorHandler};
+
+    const auto spec{test::CDataFrameAnalysisSpecificationFactory::predictionSpec(
+        "regression", "dep_var", 5, 6, 13000000, 0, 0)};
+    rapidjson::Document jsonParameters;
+    jsonParameters.Parse("{"
+                         "  \"dependent_variable\": \"dep_var\","
+                         "  \"prediction_field_name\": \"is_training\""
+                         "}");
+    const auto parameters{
+        api::CDataFrameTrainBoostedTreeRegressionRunner::parameterReader().read(jsonParameters)};
+    api::CDataFrameTrainBoostedTreeRegressionRunner runner(*spec, parameters);
+
+    BOOST_TEST_REQUIRE(errors.size() == 1);
+    BOOST_TEST_REQUIRE(errors[0] == "Input error: prediction_field_name must not be equal to any of [is_training].");
+}
+
+BOOST_AUTO_TEST_SUITE_END()
diff --git a/lib/api/unittest/Makefile b/lib/api/unittest/Makefile
@@ -31,6 +31,7 @@ SRCS=\
 	CDataFrameAnalyzerOutlierTest.cc \
 	CDataFrameAnalyzerTrainingTest.cc \
 	CDataFrameTrainBoostedTreeClassifierRunnerTest.cc \
+	CDataFrameTrainBoostedTreeRegressionRunnerTest.cc \
 	CDataFrameMockAnalysisRunner.cc \
 	CDetectionRulesJsonParserTest.cc \
 	CFieldConfigTest.cc \