From 97cbe0c6e2973c1f74ef95e0400ec45fe8b2229a Mon Sep 17 00:00:00 2001
From: Earthson Lu <Earthson.Lu@gmail.com>
Date: Mon, 11 Jan 2016 19:05:32 +0800
Subject: [PATCH] add checkSchema for union datatype

---
 .../spark/ml/feature/CountVectorizer.scala      |  3 ++-
 .../org/apache/spark/ml/util/SchemaUtils.scala  | 17 +++++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
index 10dcda2382f4..48c5128b6581 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
@@ -71,7 +71,8 @@ private[feature] trait CountVectorizerParams extends Params with HasInputCol wit
   /** Validates and transforms the input schema. */
   protected def validateAndTransformSchema(schema: StructType): StructType = {
     validateParams()
-    SchemaUtils.checkColumnType(schema, $(inputCol), new ArrayType(StringType, true))
+    val typeCandidates = List(new ArrayType(StringType, true), new ArrayType(StringType, false))
+    SchemaUtils.checkUnionColumnType(schema, $(inputCol), typeCandidates)
     SchemaUtils.appendColumn(schema, $(outputCol), new VectorUDT)
   }
 
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
index 76f651488aef..eddf076df5bf 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/util/SchemaUtils.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/util/SchemaUtils.scala
@@ -43,6 +43,23 @@ private[spark] object SchemaUtils {
       s"Column $colName must be of type $dataType but was actually $actualDataType.$message")
   }
 
+  /**
+    * Check whether the given column is one of the required data types.
+    * @param colName  column name
+    * @param dataTypes  required column data type
+    */
+  def checkUnionColumnType(
+                       schema: StructType,
+                       colName: String,
+                       dataTypes: Seq[DataType],
+                       msg: String = ""): Unit = {
+    val actualDataType = schema(colName).dataType
+    val message = if (msg != null && msg.trim.length > 0) " " + msg else ""
+    lazy val typesMsg = dataTypes.mkString(",")
+    require(dataTypes.exists(_ equals actualDataType),
+      s"Column $colName must be one of type $typesMsg but was actually $actualDataType.$message")
+  }
+
   /**
    * Appends a new column to the input schema. This fails if the given output column already exists.
    * @param schema input schema