Fixed an issue with OneHotEncoderOp when computing categorySizes

The spark feature org.apache.spark.ml.feature.OneHotEncoderModel has two mixins for the input columns: inputCol and inputCols. We need to check which param is set and use that correct one to compute categorySizes.

Author: ltrottier-yelp

mleap-spark/src/main/scala/org/apache/spark/ml/bundle/ops/feature/OneHotEncoderOp.scala

@@ -42,14 +42,16 @@ class OneHotEncoderOp extends SimpleSparkOp[OneHotEncoderModel] {
     override def store(model: Model, obj: OneHotEncoderModel)
                       (implicit context: BundleContext[SparkBundleContext]): Model = {
       assert(context.context.dataset.isDefined, BundleHelper.sampleDataframeMessage(klazz))
-
-      val df = context.context.dataset.get
-      val categorySizes = obj.getInputCols.map { f ⇒ OneHotEncoderOp.sizeForField(df.schema(f)) }
-
+      assert(!(obj.isSet(obj.inputCol) && obj.isSet(obj.inputCols)), "OneHotEncoderModel cannot have both inputCol and inputCols set")
+      val inputCols = obj.isSet(obj.inputCol) match {
+        case true => Array(obj.getInputCol)
+        case false => obj.getInputCols
+      }
+      val df = context.context.dataset.get.schema(f)
+      val categorySizes = inputCols.map { f ⇒ OneHotEncoderOp.sizeForField(df.schema(f)) }
       model.withValue("category_sizes", Value.intList(categorySizes))
         .withValue("drop_last", Value.boolean(obj.getDropLast))
         .withValue("handle_invalid", Value.string(obj.getHandleInvalid))
-
     }
 
     override def load(model: Model)