fix dimensions in TopNOneHotEncoder

spotify · Apr 10, 2018 · cad761a · cad761a
1 parent 87a82cd
commit cad761a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/core/src/main/scala/com/spotify/featran/transformers/TopNOneHotEncoder.scala b/core/src/main/scala/com/spotify/featran/transformers/TopNOneHotEncoder.scala
@@ -83,7 +83,8 @@ private class TopNOneHotEncoder(name: String,
         b.result()
       }
 
-  override def featureDimension(c: SortedMap[String, Int]): Int = c.size
+  override def featureDimension(c: SortedMap[String, Int]): Int =
+    if (encodeMissingValue) c.size + 1 else c.size
 
   override def featureNames(c: SortedMap[String, Int]): Seq[String] = {
     val names = c.map(name + '_' + _._1)(scala.collection.breakOut)