[SPARK-13568] [ML] Create feature transformer to impute missing values

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

@@ -0,0 +1,260 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.ml.feature
+
+import org.apache.hadoop.fs.Path
+
+import org.apache.spark.SparkException
+import org.apache.spark.annotation.{Experimental, Since}
+import org.apache.spark.ml.{Estimator, Model}
+import org.apache.spark.ml.param._
+import org.apache.spark.ml.param.shared.{HasInputCols, HasOutputCol}
+import org.apache.spark.ml.util._
+import org.apache.spark.sql.{DataFrame, Dataset, Row}
+import org.apache.spark.sql.functions._
+import org.apache.spark.sql.types._
+
+/**
+ * Params for [[Imputer]] and [[ImputerModel]].
+ */
+private[feature] trait ImputerParams extends Params with HasInputCols with HasOutputCol {
+
+  /**
+   * The imputation strategy.
+   * If "mean", then replace missing values using the mean value of the feature.
+   * If "median", then replace missing values using the approximate median value of the feature.
+   * Default: mean
+   *
+   * @group param
+   */
+  final val strategy: Param[String] = new Param(this, "strategy", "strategy for imputation. " +
+    "If mean, then replace missing values using the mean value of the feature. " +
+    "If median, then replace missing values using the median value of the feature.",
+    ParamValidators.inArray[String](Imputer.supportedStrategyNames.toArray))
+
+  /** @group getParam */
+  def getStrategy: String = $(strategy)
+
+  /**
+   * The placeholder for the missing values. All occurrences of missingValue will be imputed.
+   * Note that null values are always treated as missing.
+   * Default: Double.NaN
+   *
+   * @group param
+   */
+  final val missingValue: DoubleParam = new DoubleParam(this, "missingValue",
+    "The placeholder for the missing values. All occurrences of missingValue will be imputed")
+
+  /** @group getParam */
+  def getMissingValue: Double = $(missingValue)
+
+    /**
+   * Param for output column names.
+   * @group param
+   */
+  final val outputCols: StringArrayParam = new StringArrayParam(this, "outputCols",
+    "output column names")
+
+  /** @group getParam */
+  final def getOutputCols: Array[String] = $(outputCols)
+
+  /** Validates and transforms the input schema. */
+  protected def validateAndTransformSchema(schema: StructType): StructType = {
+    require($(inputCols).length == $(outputCols).length, "inputCols and outputCols should have" +
+      "the same length")
+    val localInputCols = $(inputCols)
+    val localOutputCols = $(outputCols)
+    var outputSchema = schema
+
+    $(inputCols).indices.foreach { i =>
+      val inputCol = localInputCols(i)
+      val outputCol = localOutputCols(i)
+      val inputType = schema(inputCol).dataType
+      SchemaUtils.checkColumnTypes(schema, inputCol, Seq(DoubleType, FloatType))
+      outputSchema = SchemaUtils.appendColumn(outputSchema, outputCol, inputType)
+    }
+    outputSchema
+  }
+}
+
+/**
+ * :: Experimental ::
+ * Imputation estimator for completing missing values, either using the mean or the median
+ * of the column in which the missing values are located. The input column should be of
+ * DoubleType or FloatType. Currently Imputer does not support categorical features yet
+ * (SPARK-15041) and possibly creates incorrect values for a categorical feature.
+ *
+ * Note that the mean/median value is computed after filtering out missing values.
+ * All Null values in the input column are treated as missing, and so are also imputed.
+ */
+@Experimental
+class Imputer @Since("2.1.0")(override val uid: String)
+  extends Estimator[ImputerModel] with ImputerParams with DefaultParamsWritable {
+
+  @Since("2.1.0")
+  def this() = this(Identifiable.randomUID("imputer"))
+
+  /** @group setParam */
+  @Since("2.1.0")
+  def setInputCols(value: Array[String]): this.type = set(inputCols, value)
+
+  /** @group setParam */
+  @Since("2.1.0")
+  def setOutputCols(value: Array[String]): this.type = set(outputCols, value)
+
+  /**
+   * Imputation strategy. Available options are ["mean", "median"].
+   * @group setParam
+   */
+  @Since("2.1.0")
+  def setStrategy(value: String): this.type = set(strategy, value)
+
+  /** @group setParam */
+  @Since("2.1.0")
+  def setMissingValue(value: Double): this.type = set(missingValue, value)
+
+  setDefault(strategy -> "mean", missingValue -> Double.NaN)
+
+  override def fit(dataset: Dataset[_]): ImputerModel = {
+    transformSchema(dataset.schema, logging = true)
+    val surrogates = $(inputCols).map { inputCol =>
+      val ic = col(inputCol)
+      val filtered = dataset.select(ic.cast(DoubleType))
+        .filter(ic.isNotNull && ic =!= $(missingValue))
+        .filter(!ic.isNaN)
+      if(filtered.rdd.isEmpty()) {
+        throw new SparkException(s"surrogate cannot be computed. " +
+          s"All the values in ${inputCol} are Null, Nan or missingValue ($missingValue)")
+      }
+      val surrogate = $(strategy) match {
+        case "mean" => filtered.select(avg(inputCol)).first().getDouble(0)
+        case "median" => filtered.stat.approxQuantile(inputCol, Array(0.5), 0.001)(0)
+      }
+      surrogate.asInstanceOf[Double]
+    }
+
+    import dataset.sparkSession.implicits._
+    val surrogateDF = Seq(surrogates).toDF("surrogates")
+    copyValues(new ImputerModel(uid, surrogateDF).setParent(this))
+  }
+
+  override def transformSchema(schema: StructType): StructType = {
+    validateAndTransformSchema(schema)
+  }
+
+  override def copy(extra: ParamMap): Imputer = defaultCopy(extra)
+}
+
+@Since("2.1.0")
+object Imputer extends DefaultParamsReadable[Imputer] {
+
+  /** Set of strategy names that Imputer currently supports. */
+  private[ml] val supportedStrategyNames = Set("mean", "median")
+
+  @Since("2.1.0")
+  override def load(path: String): Imputer = super.load(path)
+}
+
+/**
+ * :: Experimental ::
+ * Model fitted by [[Imputer]].
+ *
+ * @param surrogateDF Value by which missing values in the input columns will be replaced. This
+ *    is stored using DataFrame with input column names and the corresponding surrogates.
+ */
+@Experimental
+class ImputerModel private[ml](
+    override val uid: String,
+    val surrogateDF: DataFrame)
+  extends Model[ImputerModel] with ImputerParams with MLWritable {
+
+  import ImputerModel._
+
+  /** @group setParam */
+  def setInputCols(value: Array[String]): this.type = set(inputCols, value)
+
+  /** @group setParam */
+  def setOutputCols(value: Array[String]): this.type = set(outputCols, value)
+
+  override def transform(dataset: Dataset[_]): DataFrame = {
+    transformSchema(dataset.schema, logging = true)
+    val localInputCols = $(inputCols)
+    val localOutputCols = $(outputCols)
+    var outputDF = dataset
+    val surrogates = surrogateDF.head().getSeq[Double](0)
+
+    $(inputCols).indices.foreach { i =>
+      val inputCol = localInputCols(i)
+      val outputCol = localOutputCols(i)
+      val inputType = dataset.schema(inputCol).dataType
+      val ic = col(inputCol)
+      val icSurrogate = surrogates(i)
+      outputDF = outputDF.withColumn(outputCol, when(ic.isNull, icSurrogate)
+        .when(ic === $(missingValue), icSurrogate)
+        .otherwise(ic)
+        .cast(inputType))
+    }
+    outputDF.toDF()
+  }
+
+  override def transformSchema(schema: StructType): StructType = {
+    validateAndTransformSchema(schema)
+  }
+
+  override def copy(extra: ParamMap): ImputerModel = {
+    val copied = new ImputerModel(uid, surrogateDF)
+    copyValues(copied, extra).setParent(parent)
+  }
+
+  @Since("2.1.0")
+  override def write: MLWriter = new ImputerModelWriter(this)
+}
+
+
+@Since("2.1.0")
+object ImputerModel extends MLReadable[ImputerModel] {
+
+  private[ImputerModel] class ImputerModelWriter(instance: ImputerModel) extends MLWriter {
+
+    override protected def saveImpl(path: String): Unit = {
+      DefaultParamsWriter.saveMetadata(instance, path, sc)
+      val dataPath = new Path(path, "data").toString
+      instance.surrogateDF.repartition(1).write.parquet(dataPath)
+    }
+  }
+
+  private class ImputerReader extends MLReader[ImputerModel] {
+
+    private val className = classOf[ImputerModel].getName
+
+    override def load(path: String): ImputerModel = {
+      val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
+      val dataPath = new Path(path, "data").toString
+      val surrogateDF = sqlContext.read.parquet(dataPath)
+      val model = new ImputerModel(metadata.uid, surrogateDF)
+      DefaultParamsReader.getAndSetParams(model, metadata)
+      model
+    }
+  }
+
+  @Since("2.1.0")
+  override def read: MLReader[ImputerModel] = new ImputerReader
+
+  @Since("2.1.0")
+  override def load(path: String): ImputerModel = super.load(path)
+}

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

@@ -0,0 +1,136 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.ml.feature
+
+import org.apache.spark.{SparkException, SparkFunSuite}
+import org.apache.spark.ml.util.DefaultReadWriteTest
+import org.apache.spark.mllib.util.MLlibTestSparkContext
+import org.apache.spark.mllib.util.TestingUtils._
+import org.apache.spark.sql.{DataFrame, Row}
+
+class ImputerSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
+
+  test("Imputer for Double with default missing Value NaN") {
+    val df = spark.createDataFrame( Seq(
+      (0, 1.0, 1.0, 1.0),
+      (1, 1.0, 1.0, 1.0),
+      (2, 3.0, 3.0, 3.0),
+      (3, 4.0, 4.0, 4.0),
+      (4, Double.NaN, 2.25, 1.0)
+    )).toDF("id", "value", "expected_mean", "expected_median")
+    val imputer = new Imputer().setInputCols(Array("value")).setOutputCols(Array("out"))
+    ImputerSuite.iterateStrategyTest(imputer, df)
+  }
+
+  test("Imputer should handle NaNs when computing surrogate value, if missingValue is not NaN") {
+    val df = spark.createDataFrame( Seq(
+      (0, 1.0, 1.0, 1.0),
+      (1, 3.0, 3.0, 3.0),
+      (2, Double.NaN, Double.NaN, Double.NaN),
+      (3, -1.0, 2.0, 3.0)
+    )).toDF("id", "value", "expected_mean", "expected_median")
+    val imputer = new Imputer().setInputCols(Array("value")).setOutputCols(Array("out"))
+      .setMissingValue(-1.0)
+    ImputerSuite.iterateStrategyTest(imputer, df)
+  }
+
+  test("Imputer for Float with missing Value -1.0") {
+    val df = spark.createDataFrame( Seq(
+      (0, 1.0F, 1.0F, 1.0F),
+      (1, 3.0F, 3.0F, 3.0F),
+      (2, 10.0F, 10.0F, 10.0F),
+      (3, 10.0F, 10.0F, 10.0F),
+      (4, -1.0F, 6.0F, 3.0F)
+    )).toDF("id", "value", "expected_mean", "expected_median")
+    val imputer = new Imputer().setInputCols(Array("value")).setOutputCols(Array("out"))
+      .setMissingValue(-1)
+    ImputerSuite.iterateStrategyTest(imputer, df)
+  }
+
+  test("Imputer should impute null as well as 'missingValue'") {
+    val df = spark.createDataFrame( Seq(
+      (0, 4.0, 4.0, 4.0),
+      (1, 10.0, 10.0, 10.0),
+      (2, 10.0, 10.0, 10.0),
+      (3, Double.NaN, 8.0, 10.0),
+      (4, -1.0, 8.0, 10.0)
+    )).toDF("id", "value", "expected_mean", "expected_median")
+    val df2 = df.selectExpr("*", "IF(value=-1.0, null, value) as nullable_value")
+    val imputer = new Imputer().setInputCols(Array("nullable_value")).setOutputCols(Array("out"))
+    ImputerSuite.iterateStrategyTest(imputer, df2)
+  }
+
+
+  test("Imputer throws exception when surrogate cannot be computed") {
+    val df = spark.createDataFrame( Seq(
+      (0, Double.NaN, 1.0, 1.0),
+      (1, Double.NaN, 3.0, 3.0),
+      (2, Double.NaN, Double.NaN, Double.NaN)
+    )).toDF("id", "value", "expected_mean", "expected_median")
+    Seq("mean", "median").foreach { strategy =>
+      val imputer = new Imputer().setInputCols(Array("value")).setOutputCols(Array("out"))
+        .setStrategy(strategy)
+      intercept[SparkException] {
+        val model = imputer.fit(df)
+      }
+    }
+  }
+
+  test("Imputer read/write") {
+    val t = new Imputer()
+      .setInputCols(Array("myInputCol"))
+      .setOutputCols(Array("myOutputCol"))
+      .setMissingValue(-1.0)
+    testDefaultReadWrite(t)
+  }
+
+  test("ImputerModel read/write") {
+    val spark = this.spark
+    import spark.implicits._
+    val surrogateDF = Seq(1.234).toDF("myInputCol")
+
+    val instance = new ImputerModel(
+      "myImputer", surrogateDF)
+      .setInputCols(Array("myInputCol"))
+      .setOutputCols(Array("myOutputCol"))
+    val newInstance = testDefaultReadWrite(instance)
+    assert(newInstance.surrogateDF.collect() === instance.surrogateDF.collect())
+  }
+
+}
+
+object ImputerSuite{
+
+  /**
+   * Imputation strategy. Available options are ["mean", "median"].
+   * @param df DataFrame with columns "id", "value", "expected_mean", "expected_median"
+   */
+  def iterateStrategyTest(imputer: Imputer, df: DataFrame): Unit = {
+    Seq("mean", "median").foreach { strategy =>
+      imputer.setStrategy(strategy)
+      val model = imputer.fit(df)
+      model.transform(df).select("expected_" + strategy, "out").collect().foreach {
+        case Row(exp: Float, out: Float) =>
+          assert((exp.isNaN && out.isNaN) || (exp == out),
+            s"Imputed values differ. Expected: $exp, actual: $out")
+        case Row(exp: Double, out: Double) =>
+          assert((exp.isNaN && out.isNaN) || (exp ~== out absTol 1e-5),
+            s"Imputed values differ. Expected: $exp, actual: $out")
+      }
+    }
+  }
+}