[SPARK-20109][MLlib] Rewrote toBlockMatrix method on IndexedRowMatrix

mllib/src/main/scala/org/apache/spark/mllib/linalg/distributed/IndexedRowMatrix.scala

@@ -98,6 +98,7 @@ class IndexedRowMatrix @Since("1.0.0") (
     toBlockMatrix(1024, 1024)
   }
 
+
   /**
    * Converts to BlockMatrix. Creates blocks of `SparseMatrix`.
    * @param rowsPerBlock The number of rows of each block. The blocks at the bottom edge may have
@@ -112,6 +113,67 @@ class IndexedRowMatrix @Since("1.0.0") (
     toCoordinateMatrix().toBlockMatrix(rowsPerBlock, colsPerBlock)
   }
 
+  /**
+    * Converts to BlockMatrix. Creates blocks of `DenseMatrix` with size 1024 x 1024.
+    */
+  def toBlockMatrixDense(): BlockMatrix = {
+    toBlockMatrixDense(1024, 1024)
+  }
+
+  /**
+    * Converts to BlockMatrix. Creates blocks of `DenseMatrix`.
+    * @param rowsPerBlock The number of rows of each block. The blocks at the bottom edge may have
+    *                     a smaller value. Must be an integer value greater than 0.
+    * @param colsPerBlock The number of columns of each block. The blocks at the right edge may have
+    *                     a smaller value. Must be an integer value greater than 0.
+    * @return a [[BlockMatrix]]
+    */
+  def toBlockMatrixDense(rowsPerBlock: Int, colsPerBlock: Int): BlockMatrix = {
+    require(rowsPerBlock > 0,
+      s"rowsPerBlock needs to be greater than 0. rowsPerBlock: $rowsPerBlock")
+    require(colsPerBlock > 0,
+      s"colsPerBlock needs to be greater than 0. colsPerBlock: $colsPerBlock")
+
+    val m = numRows()
+    val n = numCols()
+    val lastRowBlockIndex = m / rowsPerBlock
+    val lastColBlockIndex = n / colsPerBlock
+    val lastRowBlockSize = (m % rowsPerBlock).toInt
+    val lastColBlockSize = (n % colsPerBlock).toInt
+    val numRowBlocks = math.ceil(m.toDouble / rowsPerBlock).toInt
+    val numColBlocks = math.ceil(n.toDouble / colsPerBlock).toInt
+
+    val blocks: RDD[((Int, Int), Matrix)] = rows.flatMap({ ir =>
+      val blockRow = ir.index / rowsPerBlock
+      val rowInBlock = ir.index % rowsPerBlock
+
+      ir.vector.toArray
+        .grouped(colsPerBlock)
+        .zipWithIndex
+        .map({ case (values, blockColumn) =>
+          ((blockRow.toInt, blockColumn), (rowInBlock.toInt, values))
+        })
+    }).groupByKey(GridPartitioner(numRowBlocks, numColBlocks, rowsPerBlock, colsPerBlock)).map({
+      case ((blockRow, blockColumn), itr) =>
+        val actualNumRows: Int =
+          if (blockRow == lastRowBlockIndex) lastRowBlockSize else rowsPerBlock
+        val actualNumColumns: Int =
+          if (blockColumn == lastColBlockIndex) lastColBlockSize else colsPerBlock
+
+        val arraySize = actualNumRows * actualNumColumns
+        val matrixAsArray = new Array[Double](arraySize)
+        itr.foreach({ case (rowWithinBlock, values) =>
+          var i = 0
+          while (i < values.length) {
+            matrixAsArray.update(i * actualNumRows + rowWithinBlock, values(i))
+            i += 1
+          }
+        })
+        ((blockRow, blockColumn), new DenseMatrix(actualNumRows, actualNumColumns, matrixAsArray))
+    })
+    new BlockMatrix(blocks, rowsPerBlock, colsPerBlock)
+  }
+
   /**
    * Converts this matrix to a
    * [[org.apache.spark.mllib.linalg.distributed.CoordinateMatrix]].

mllib/src/test/scala/org/apache/spark/mllib/linalg/distributed/IndexedRowMatrixSuite.scala

@@ -89,11 +89,42 @@ class IndexedRowMatrixSuite extends SparkFunSuite with MLlibTestSparkContext {
 
   test("toBlockMatrix") {
     val idxRowMat = new IndexedRowMatrix(indexedRows)
+
+    // Tests when n % colsPerBlock != 0
+    val blockMat = idxRowMat.toBlockMatrix(2, 2)
+    assert(blockMat.numRows() === m)
+    assert(blockMat.numCols() === n)
+    assert(blockMat.toBreeze() === idxRowMat.toBreeze())
+
+    // Tests when m % rowsPerBlock != 0
+    val blockMat2 = idxRowMat.toBlockMatrix(3, 1)
+    assert(blockMat2.numRows() === m)
+    assert(blockMat2.numCols() === n)
+    assert(blockMat2.toBreeze() === idxRowMat.toBreeze())
+
+    intercept[IllegalArgumentException] {
+      idxRowMat.toBlockMatrix(-1, 2)
+    }
+    intercept[IllegalArgumentException] {
+      idxRowMat.toBlockMatrix(2, 0)
+    }
+  }
+
+  test("toBlockMatrixDense") {
+    val idxRowMat = new IndexedRowMatrix(indexedRows)
+
+    // Tests when n % colsPerBlock != 0
     val blockMat = idxRowMat.toBlockMatrix(2, 2)
     assert(blockMat.numRows() === m)
     assert(blockMat.numCols() === n)
     assert(blockMat.toBreeze() === idxRowMat.toBreeze())
 
+    // Tests when m % rowsPerBlock != 0
+    val blockMat2 = idxRowMat.toBlockMatrix(3, 1)
+    assert(blockMat2.numRows() === m)
+    assert(blockMat2.numCols() === n)
+    assert(blockMat2.toBreeze() === idxRowMat.toBreeze())
+
     intercept[IllegalArgumentException] {
       idxRowMat.toBlockMatrix(-1, 2)
     }