[ML] Correct missing category handling (elastic#1042)

Author: tveasey

include/core/CDataFrame.h

@@ -20,6 +20,7 @@
 #include <cstdint>
 #include <functional>
 #include <iterator>
+#include <limits>
 #include <memory>
 #include <vector>
 
@@ -490,7 +491,9 @@ class CORE_EXPORT CDataFrame final {
                                            std::size_t numberColumns);
 
     //! Get the value to use for a missing element in a data frame.
-    static double valueOfMissing();
+    static constexpr double valueOfMissing() {
+        return std::numeric_limits<double>::quiet_NaN();
+    }
 
 private:
     using TStrSizeUMap = boost::unordered_map<std::string, std::size_t>;

include/maths/CDataFrameUtils.h

@@ -122,9 +122,15 @@ class MATHS_EXPORT CDataFrameUtils : private core::CNonInstantiatable {
         COneHotCategoricalColumnValue(std::size_t column, std::size_t category)
             : CColumnValue{column}, m_Category{category} {}
         double operator()(const TRowRef& row) const override {
+            if (isMissing(row[this->column()])) {
+                return core::CDataFrame::valueOfMissing();
+            }
             return static_cast<std::size_t>(row[this->column()]) == m_Category ? 1.0 : 0.0;
         }
         double operator()(const TFloatVec& row) const override {
+            if (isMissing(row[this->column()])) {
+                return core::CDataFrame::valueOfMissing();
+            }
             return static_cast<std::size_t>(row[this->column()]) == m_Category ? 1.0 : 0.0;
         }
         std::size_t hash() const override { return m_Category; }
@@ -140,10 +146,16 @@ class MATHS_EXPORT CDataFrameUtils : private core::CNonInstantiatable {
         CFrequencyCategoricalColumnValue(std::size_t column, const TDoubleVec& frequencies)
             : CColumnValue{column}, m_Frequencies{&frequencies} {}
         double operator()(const TRowRef& row) const override {
+            if (isMissing(row[this->column()])) {
+                return core::CDataFrame::valueOfMissing();
+            }
             std::size_t category{static_cast<std::size_t>(row[this->column()])};
             return (*m_Frequencies)[category];
         }
         double operator()(const TFloatVec& row) const override {
+            if (isMissing(row[this->column()])) {
+                return core::CDataFrame::valueOfMissing();
+            }
             std::size_t category{static_cast<std::size_t>(row[this->column()])};
             return (*m_Frequencies)[category];
         }
@@ -166,10 +178,16 @@ class MATHS_EXPORT CDataFrameUtils : private core::CNonInstantiatable {
             : CColumnValue{column}, m_RareCategories{&rareCategories}, m_TargetMeanValues{&targetMeanValues} {
         }
         double operator()(const TRowRef& row) const override {
+            if (isMissing(row[this->column()])) {
+                return core::CDataFrame::valueOfMissing();
+            }
             std::size_t category{static_cast<std::size_t>(row[this->column()])};
             return this->isRare(category) ? 0.0 : (*m_TargetMeanValues)[category];
         }
         double operator()(const TFloatVec& row) const override {
+            if (isMissing(row[this->column()])) {
+                return core::CDataFrame::valueOfMissing();
+            }
             std::size_t category{static_cast<std::size_t>(row[this->column()])};
             return this->isRare(category) ? 0.0 : (*m_TargetMeanValues)[category];
         }

lib/core/CDataFrame.cc

@@ -413,10 +413,6 @@ std::size_t CDataFrame::estimateMemoryUsage(bool inMainMemory,
     return inMainMemory ? numberRows * numberColumns * sizeof(float) : 0;
 }
 
-double CDataFrame::valueOfMissing() {
-    return std::numeric_limits<double>::quiet_NaN();
-}
-
 CDataFrame::TRowFuncVecBoolPr
 CDataFrame::parallelApplyToAllRows(std::size_t numberThreads,
                                    std::size_t beginRows,

lib/maths/CDataFrameCategoryEncoder.cc

@@ -7,6 +7,7 @@
 #include <maths/CDataFrameCategoryEncoder.h>
 
 #include <core/CContainerPrinter.h>
+#include <core/CDataFrame.h>
 #include <core/CLogger.h>
 #include <core/CPackedBitVector.h>
 #include <core/CPersistUtils.h>
@@ -452,7 +453,9 @@ EEncoding CDataFrameCategoryEncoder::COneHotEncoding::type() const {
 }
 
 double CDataFrameCategoryEncoder::COneHotEncoding::encode(double value) const {
-    return static_cast<std::size_t>(value) == m_HotCategory;
+    return CDataFrameUtils::isMissing(value)
+               ? core::CDataFrame::valueOfMissing()
+               : static_cast<std::size_t>(value) == m_HotCategory;
 }
 
 bool CDataFrameCategoryEncoder::COneHotEncoding::isBinary() const {
@@ -503,6 +506,9 @@ EEncoding CDataFrameCategoryEncoder::CMappedEncoding::type() const {
 }
 
 double CDataFrameCategoryEncoder::CMappedEncoding::encode(double value) const {
+    if (CDataFrameUtils::isMissing(value)) {
+        return core::CDataFrame::valueOfMissing();
+    }
     std::size_t category{static_cast<std::size_t>(value)};
     return category < m_Map.size() ? m_Map[category] : m_Fallback;
 }

lib/maths/CDataFrameUtils.cc

@@ -824,9 +824,11 @@ CDataFrameUtils::TSizeDoublePrVecVecVec CDataFrameUtils::categoricalMicWithColum
                 1, 0, frame.numberRows(),
                 [&](TRowItr beginRows, TRowItr endRows) {
                     for (auto row = beginRows; row != endRows; ++row) {
+                        if (isMissing((*row)[i]) || isMissing(target(*row))) {
+                            continue;
+                        }
                         std::size_t category{static_cast<std::size_t>((*row)[i])};
-                        if (frequencies[i][category] >= minimumFrequency &&
-                            isMissing(target(*row)) == false) {
+                        if (frequencies[i][category] >= minimumFrequency) {
                             sampler.sample(*row);
                         }
                     }
@@ -905,6 +907,9 @@ CDataFrameUtils::TSizeDoublePrVecVecVec CDataFrameUtils::categoricalMicWithColum
 
             encoders.clear();
             for (const auto& sample : samples) {
+                if (isMissing(sample[i])) {
+                    continue;
+                }
                 std::size_t category{static_cast<std::size_t>(sample[i])};
                 if (frequencies[i][category] >= minimumFrequency) {
                     auto encoder = makeEncoder(i, i, category);