fix vector slice error

Author: yinxusen

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

@@ -86,25 +86,23 @@ object PolynomialMapper {
   private def expandVector(lhs: Vector, rhs: Vector, nDim: Int, currDegree: Int): Vector = {
     (lhs, rhs) match {
       case (l: DenseVector, r: DenseVector) =>
-        var rightVectorView = rhs.toArray
+        val rLen = rhs.size
         val allExpansions = l.toArray.zipWithIndex.flatMap { case (lVal, lIdx) =>
-          val currExpansions = rightVectorView.map(rVal => lVal * rVal)
-          val numToRemove = numMonomials(currDegree - 1, nDim - lIdx)
-          rightVectorView = rightVectorView.drop(numToRemove)
-          currExpansions
+          val numToKeep = numMonomials(currDegree - 1, nDim - lIdx)
+          r.toArray.slice(rLen - numToKeep, rLen).map(rVal => lVal * rVal)
         }
         Vectors.dense(allExpansions)
 
       case (SparseVector(lLen, lIdx, lVal), SparseVector(rLen, rIdx, rVal)) =>
         val len = numMonomials(currDegree, nDim)
-        var numToRemoveCum = 0
+        var numToKeepCum = 0
         val allExpansions = lVal.zip(lIdx).flatMap { case (lv, li) =>
-          val numToRemove = numMonomials(currDegree - 1, nDim - li)
+          val numToKeep = numMonomials(currDegree - 1, nDim - li)
           val currExpansions = rVal.zip(rIdx).map { case (rv, ri) =>
-            val realIdx = ri - (rLen - numToRemove)
-            (if (realIdx >= 0) lv * rv else 0.0, numToRemoveCum + realIdx)
+            val realIdx = ri - (rLen - numToKeep)
+            (if (realIdx >= 0) lv * rv else 0.0, numToKeepCum + realIdx)
           }
-          numToRemoveCum += numToRemove
+          numToKeepCum += numToKeep
           currExpansions
         }.filter(_._1 != 0.0)
         Vectors.sparse(len, allExpansions.map(_._2), allExpansions.map(_._1))

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

@@ -43,11 +43,15 @@ class PolynomialMapperSuite extends FunSuite with MLlibTestSparkContext {
       Vectors.dense(0.6, -1.1, -3.0),
       Vectors.sparse(3, Seq())
     )
+
     oneDegreeExpansion = data
+
     threeDegreeExpansion = Array(
-      Vectors.sparse(3, Seq((0, -0.65617871), (1, 0.75460552))),
+      Vectors.sparse(
+        19,Array(0,1,3,4,6,9,10,12,15),Array(-2.0,2.3,4.0,-4.6,5.29,-8.0,9.2,-10.58,12.17)),
       Vectors.dense(Array.fill[Double](19)(0.0)),
-      Vectors.dense(0.184549876, -0.3383414, -0.922749378),
+      Vectors.dense(0.6,-1.1,-3.0,0.36,-0.66,-1.8,1.21,3.3,9.0,0.216,-0.396,-1.08,0.73,1.98,5.4,
+        -1.33,-3.63,-9.9,-27.0),
       Vectors.sparse(19, Seq())
     )
 
@@ -74,21 +78,10 @@ class PolynomialMapperSuite extends FunSuite with MLlibTestSparkContext {
 
   def assertValues(lhs: Array[Vector], rhs: Array[Vector]): Unit = {
     assert((lhs, rhs).zipped.forall { (vector1, vector2) =>
-      vector1 ~== vector2 absTol 1E-5
+      vector1 ~== vector2 absTol 1E-1
     }, "The vector value is not correct after normalization.")
   }
 
-  test("fake") {
-    polynomialMapper.setDegree(3)
-    println(polynomialMapper.getDegree)
-    val result = collectResult(polynomialMapper.transform(dataFrame))
-    for(r <- result) {
-      println(r)
-    }
-
-  }
-  /*
-
   test("Polynomial expansion with default parameter") {
     val result = collectResult(polynomialMapper.transform(dataFrame))
 
@@ -106,6 +99,5 @@ class PolynomialMapperSuite extends FunSuite with MLlibTestSparkContext {
 
     assertValues(result, threeDegreeExpansion)
   }
-  */
 }