[SPARK-12711][ML] ML StopWordsRemover does not protect itself from co…

…lumn name duplication Fixes problem and verifies fix by test suite. Also - adds optional parameter: nullable (Boolean) to: SchemaUtils.appendColumn and deduplicates SchemaUtils.appendColumn functions. Author: Grzegorz Chilkiewicz <[email protected]> Closes apache#10741 from grzegorz-chilkiewicz/master.
leonelfolmer · Feb 2, 2016 · b1835d7 · b1835d7
1 parent 358300c
commit b1835d7
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 8 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/StopWordsRemover.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/StopWordsRemover.scala
@@ -149,9 +149,7 @@ class StopWordsRemover(override val uid: String)
     val inputType = schema($(inputCol)).dataType
     require(inputType.sameType(ArrayType(StringType)),
       s"Input type must be ArrayType(StringType) but got $inputType.")
-    val outputFields = schema.fields :+
-      StructField($(outputCol), inputType, schema($(inputCol)).nullable)
-    StructType(outputFields)
+    SchemaUtils.appendColumn(schema, $(outputCol), inputType, schema($(inputCol)).nullable)
   }
 
   override def copy(extra: ParamMap): StopWordsRemover = defaultCopy(extra)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/util/SchemaUtils.scala b/mllib/src/main/scala/org/apache/spark/ml/util/SchemaUtils.scala
@@ -71,12 +71,10 @@ private[spark] object SchemaUtils {
   def appendColumn(
       schema: StructType,
       colName: String,
-      dataType: DataType): StructType = {
+      dataType: DataType,
+      nullable: Boolean = false): StructType = {
     if (colName.isEmpty) return schema
-    val fieldNames = schema.fieldNames
-    require(!fieldNames.contains(colName), s"Column $colName already exists.")
-    val outputFields = schema.fields :+ StructField(colName, dataType, nullable = false)
-    StructType(outputFields)
+    appendColumn(schema, StructField(colName, dataType, nullable))
   }
 
   /**

diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/StopWordsRemoverSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/StopWordsRemoverSuite.scala
@@ -89,4 +89,19 @@ class StopWordsRemoverSuite
       .setCaseSensitive(true)
     testDefaultReadWrite(t)
   }
+
+  test("StopWordsRemover output column already exists") {
+    val outputCol = "expected"
+    val remover = new StopWordsRemover()
+      .setInputCol("raw")
+      .setOutputCol(outputCol)
+    val dataSet = sqlContext.createDataFrame(Seq(
+      (Seq("The", "the", "swift"), Seq("swift"))
+    )).toDF("raw", outputCol)
+
+    val thrown = intercept[IllegalArgumentException] {
+      testStopWordsRemover(remover, dataSet)
+    }
+    assert(thrown.getMessage == s"requirement failed: Column $outputCol already exists.")
+  }
 }