ensure the functionality

Author: yinxusen

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

@@ -62,7 +62,7 @@ private[feature] trait Word2VecBase extends Params
    */
   final val seed = new LongParam(this, "seed", "a random seed to random an initial vector")
 
-  setDefault(seed -> Utils.random.nextLong())
+  setDefault(seed -> 42L)
 
   /** @group getParam */
   def getSeed: Long = getOrDefault(seed)
@@ -77,12 +77,15 @@ private[feature] trait Word2VecBase extends Params
   /** @group getParam */
   def getMinCount: Int = getOrDefault(minCount)
 
+  setDefault(stepSize -> 0.025)
+  setDefault(maxIter -> 1)
+
   /**
    * Validate and transform the input schema.
    */
   protected def validateAndTransformSchema(schema: StructType, paramMap: ParamMap): StructType = {
     val map = extractParamMap(paramMap)
-    SchemaUtils.checkColumnType(schema, map(inputCol), new ArrayType(new StringType, false))
+    SchemaUtils.checkColumnType(schema, map(inputCol), new ArrayType(StringType, true))
     SchemaUtils.appendColumn(schema, map(outputCol), new VectorUDT)
   }
 }
@@ -166,8 +169,14 @@ class Word2VecModel private[ml] (
     val map = extractParamMap(paramMap)
     val bWordVectors = dataset.sqlContext.sparkContext.broadcast(wordVectors)
     val word2Vec = udf { v: Seq[String] =>
-      v.map(bWordVectors.value.transform).foldLeft(Vectors.zeros(map(vectorSize))) { (cum, vec) =>
-        Vectors.dense(cum.toArray.zip(vec.toArray).map(x => x._1 + x._2))
+      if (v.size == 0) {
+        Vectors.zeros(map(vectorSize))
+      } else {
+        Vectors.dense(
+          v.map(bWordVectors.value.getVectors).foldLeft(Array.fill[Double](map(vectorSize))(0)) {
+            (cum, vec) => cum.zip(vec).map(x => x._1 + x._2)
+          }.map(_ / v.size)
+        )
       }
     }
     dataset.withColumn(map(outputCol), word2Vec(col(map(inputCol))))

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

@@ -31,57 +31,34 @@ class Word2VecSuite extends FunSuite with MLlibTestSparkContext {
     import sqlContext.implicits._
 
     val sentence = "a b " * 100 + "a c " * 10
-    val localDoc = Seq(sentence, sentence)
-    val doc = sc.parallelize(localDoc)
-      .map(line => line.split(" "))
-    val docDF = doc.map(text => Tuple1(text)).toDF("text")
+    val numOfWords = sentence.split(" ").size
+    val doc = sc.parallelize(Seq(sentence, sentence)).map(line => line.split(" "))
 
-    val model = new Word2Vec()
-      .setVectorSize(3)
-      .setSeed(42L)
-      .setInputCol("text")
-      .setMaxIter(1)
-      .fit(docDF)
-
-    val words = sc.parallelize(Seq("a", "b", "c"))
     val codes = Map(
-      "a" -> Vectors.dense(-0.2811822295188904,-0.6356269121170044,-0.3020961284637451),
-      "b" -> Vectors.dense(1.0309048891067505,-1.29472815990448,0.22276712954044342),
-      "c" -> Vectors.dense(-0.08456747233867645,0.5137411952018738,0.11731560528278351)
+      "a" -> Array(-0.2811822295188904,-0.6356269121170044,-0.3020961284637451),
+      "b" -> Array(1.0309048891067505,-1.29472815990448,0.22276712954044342),
+      "c" -> Array(-0.08456747233867645,0.5137411952018738,0.11731560528278351)
     )
 
-    val synonyms = Map(
-      "a" -> Map("b" -> 0.3680490553379059),
-      "b" -> Map("a" -> 0.3680490553379059),
-      "c" -> Map("b" -> -0.8148014545440674)
-    )
-    val wordsDF = words.map(word => Tuple3(word, codes(word), synonyms(word)))
-      .toDF("word", "realCode", "realSynonyms")
+    val expected = doc.map { sentence =>
+      Vectors.dense(sentence.map(codes.apply).reduce((word1, word2) =>
+        word1.zip(word2).map { case (v1, v2) => v1 + v2 }
+      ).map(_ / numOfWords))
+    }
 
-    val res = model
-      .setInputCol("word")
-      .setCodeCol("code")
-      .setSynonymsCol("syn")
-      .setNumSynonyms(1)
-      .transform(wordsDF)
+    val docDF = doc.zip(expected).toDF("text", "expected")
 
-    assert(
-      res.select("code", "realCode")
-        .map { case Row(c: Vector, rc: Vector) => (c, rc) }
-        .collect()
-        .forall { case (vector1, vector2) =>
-          vector1 ~== vector2 absTol 1E-5
-        }, "The code is not correct after transforming."
-    )
+    val model = new Word2Vec()
+      .setVectorSize(3)
+      .setInputCol("text")
+      .setOutputCol("result")
+      .fit(docDF)
+
+    model.transform(docDF).select("result", "expected").collect().foreach {
+      case Row(vector1: Vector, vector2: Vector) =>
+        assert(vector1 ~== vector2 absTol 1E-5, "Transformed vector is different with expected.")
+    }
 
-    assert(
-      res.select("syn", "realSynonyms")
-        .map { case Row(s: Map[String, Double], rs: Map[String, Double]) => (s, rs) }
-        .collect()
-        .forall { case (map1, map2) =>
-          map1 == map2
-        }, "The synonyms are not correct after transforming."
-    )
   }
 }