refine IDF transformer with new interfaces

Author: yinxusen

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

@@ -20,26 +20,29 @@ package org.apache.spark.ml.feature
 import org.apache.spark.annotation.AlphaComponent
 import org.apache.spark.ml._
 import org.apache.spark.ml.param._
+import org.apache.spark.ml.param.shared._
+import org.apache.spark.ml.util.SchemaUtils
 import org.apache.spark.mllib.feature
 import org.apache.spark.mllib.linalg.{Vector, VectorUDT}
 import org.apache.spark.sql._
 import org.apache.spark.sql.functions._
-import org.apache.spark.sql.types.{StructField, StructType}
+import org.apache.spark.sql.types.StructType
 
 /**
  * Params for [[IDF]] and [[IDFModel]].
  */
-private[feature] trait IDFParams extends Params with HasInputCol with HasOutputCol {
+private[feature] trait IDFBase extends Params with HasInputCol with HasOutputCol {
 
   /**
    * The minimum of documents in which a term should appear.
    * @group param
    */
   val minDocFreq = new IntParam(
-    this, "minDocFreq", "minimum of documents in which a term should appear for filtering", Some(0))
+    this, "minDocFreq", "minimum of documents in which a term should appear for filtering")
+  setDefault(minDocFreq -> 0)
 
   /** @group getParam */
-  def getMinDocFreq: Int = get(minDocFreq)
+  def getMinDocFreq: Int = getOrDefault(minDocFreq)
 
   /** @group setParam */
   def setMinDocFreq(value: Int): this.type = set(minDocFreq, value)
@@ -48,14 +51,9 @@ private[feature] trait IDFParams extends Params with HasInputCol with HasOutputC
    * Validate and transform the input schema.
    */
   protected def validateAndTransformSchema(schema: StructType, paramMap: ParamMap): StructType = {
-    val map = this.paramMap ++ paramMap
-    val inputType = schema(map(inputCol)).dataType
-    require(inputType.isInstanceOf[VectorUDT],
-      s"Input column ${map(inputCol)} must be a vector column")
-    require(!schema.fieldNames.contains(map(outputCol)),
-      s"Output column ${map(outputCol)} already exists.")
-    val outputFields = schema.fields :+ StructField(map(outputCol), new VectorUDT, false)
-    StructType(outputFields)
+    val map = extractParamMap(paramMap)
+    SchemaUtils.checkColumnType(schema, map(inputCol), new VectorUDT)
+    SchemaUtils.appendColumn(schema, map(outputCol), new VectorUDT)
   }
 }
 
@@ -64,7 +62,7 @@ private[feature] trait IDFParams extends Params with HasInputCol with HasOutputC
  * Compute the Inverse Document Frequency (IDF) given a collection of documents.
  */
 @AlphaComponent
-class IDF extends Estimator[IDFModel] with IDFParams {
+class IDF extends Estimator[IDFModel] with IDFBase {
 
   /** @group setParam */
   def setInputCol(value: String): this.type = set(inputCol, value)
@@ -74,7 +72,7 @@ class IDF extends Estimator[IDFModel] with IDFParams {
 
   override def fit(dataset: DataFrame, paramMap: ParamMap): IDFModel = {
     transformSchema(dataset.schema, paramMap, logging = true)
-    val map = this.paramMap ++ paramMap
+    val map = extractParamMap(paramMap)
     val input = dataset.select(map(inputCol)).map { case Row(v: Vector) => v }
     val idf = new feature.IDF(getMinDocFreq).fit(input)
     val model = new IDFModel(this, map, idf)
@@ -96,7 +94,7 @@ class IDFModel private[ml] (
     override val parent: IDF,
     override val fittingParamMap: ParamMap,
     idfModel: feature.IDFModel)
-  extends Model[IDFModel] with IDFParams {
+  extends Model[IDFModel] with IDFBase {
 
   /** @group setParam */
   def setInputCol(value: String): this.type = set(inputCol, value)
@@ -106,9 +104,9 @@ class IDFModel private[ml] (
 
   override def transform(dataset: DataFrame, paramMap: ParamMap): DataFrame = {
     transformSchema(dataset.schema, paramMap, logging = true)
-    val map = this.paramMap ++ paramMap
-    val idf: Vector => Vector = (vec) => idfModel.transform(vec)
-    dataset.withColumn(map(outputCol), callUDF(idf, new VectorUDT, col(map(inputCol))))
+    val map = extractParamMap(paramMap)
+    val idf = udf { vec: Vector => idfModel.transform(vec) }
+    dataset.withColumn(map(outputCol), idf(col(map(inputCol))))
   }
 
   override def transformSchema(schema: StructType, paramMap: ParamMap): StructType = {