And switch to using handleInvalid instead of skipInvalid

Author: holdenk

mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala

@@ -33,7 +33,7 @@ import org.apache.spark.util.collection.OpenHashMap
  * Base trait for [[StringIndexer]] and [[StringIndexerModel]].
  */
 private[feature] trait StringIndexerBase extends Params with HasInputCol with HasOutputCol
-    with HasSkipInvalid {
+    with HasHandleInvalid {
 
   /** Validates and transforms the input schema. */
   protected def validateAndTransformSchema(schema: StructType): StructType = {
@@ -66,8 +66,8 @@ class StringIndexer(override val uid: String) extends Estimator[StringIndexerMod
   def this() = this(Identifiable.randomUID("strIdx"))
 
   /** @group setParam */
-  def setSkipInvalid(value: Boolean): this.type = set(skipInvalid, value)
-  setDefault(skipInvalid, false)
+  def setHandleInvalid(value: String): this.type = set(handleInvalid, value)
+  setDefault(handleInvalid, "error")
 
   /** @group setParam */
   def setInputCol(value: String): this.type = set(inputCol, value)
@@ -115,8 +115,8 @@ class StringIndexerModel private[ml] (
   }
 
   /** @group setParam */
-  def setSkipInvalid(value: Boolean): this.type = set(skipInvalid, value)
-  setDefault(skipInvalid, false)
+  def setHandleInvalid(value: String): this.type = set(handleInvalid, value)
+  setDefault(handleInvalid, "error")
 
   /** @group setParam */
   def setInputCol(value: String): this.type = set(inputCol, value)
@@ -143,13 +143,14 @@ class StringIndexerModel private[ml] (
     val metadata = NominalAttribute.defaultAttr
       .withName(outputColName).withValues(labels).toMetadata()
     // If we are skipping invalid records, filter them out.
-    val filteredDataset = if (getSkipInvalid) {
-      val filterer = udf { label: String =>
-        labelToIndex.contains(label)
+    val filteredDataset = (getHandleInvalid) match {
+      case "skip" => {
+        val filterer = udf { label: String =>
+          labelToIndex.contains(label)
+        }
+        dataset.where(filterer(dataset($(inputCol))))
       }
-      dataset.where(filterer(dataset($(inputCol))))
-    } else {
-      dataset
+      case _ =>  dataset
     }
     filteredDataset.select(col("*"),
       indexer(dataset($(inputCol)).cast(StringType)).as(outputColName, metadata))

mllib/src/main/scala/org/apache/spark/ml/param/shared/SharedParamsCodeGen.scala

@@ -54,7 +54,7 @@ private[shared] object SharedParamsCodeGen {
         isValid = "ParamValidators.gtEq(1)"),
       ParamDesc[Boolean]("fitIntercept", "whether to fit an intercept term", Some("true")),
       ParamDesc[String]("handleInvalid", "how to handle invalid entries",
-        isValid = "ParamValidators.inArray(List(\"skip\", \"error\"))"),
+        isValid = "ParamValidators.inArray(Array(\"skip\", \"error\"))"),
       ParamDesc[Boolean]("standardization", "whether to standardize the training features" +
         " before fitting the model.", Some("true")),
       ParamDesc[Long]("seed", "random seed", Some("this.getClass.getName.hashCode.toLong")),

mllib/src/main/scala/org/apache/spark/ml/param/shared/sharedParams.scala

@@ -241,7 +241,7 @@ private[ml] trait HasHandleInvalid extends Params {
    * Param for how to handle invalid entries.
    * @group param
    */
-  final val handleInvalid: Param[String] = new Param[String](this, "handleInvalid", "how to handle invalid entries", ParamValidators.inArray(List("skip", "error")))
+  final val handleInvalid: Param[String] = new Param[String](this, "handleInvalid", "how to handle invalid entries", ParamValidators.inArray(Array("skip", "error")))
 
   /** @group getParam */
   final def getHandleInvalid: String = $(handleInvalid)

mllib/src/test/scala/org/apache/spark/ml/feature/StringIndexerSuite.scala

@@ -66,7 +66,7 @@ class StringIndexerSuite extends SparkFunSuite with MLlibTestSparkContext {
     val indexerSkipInvalid = new StringIndexer()
       .setInputCol("label")
       .setOutputCol("labelIndex")
-      .setSkipInvalid(true)
+      .setHandleInvalid("skip")
       .fit(df)
     // Verify that we skip the c record
     val transformed = indexerSkipInvalid.transform(df2)