From cad761a83e6ab2257857e42b676dc3016b78d5bc Mon Sep 17 00:00:00 2001 From: Neville Li Date: Tue, 10 Apr 2018 16:45:52 -0400 Subject: [PATCH] fix dimensions in TopNOneHotEncoder --- .../com/spotify/featran/transformers/TopNOneHotEncoder.scala | 3 ++- 1 file changed, 2 insertions(+), 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 index 529b54eb..c69461f1 100644 --- 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)