added better comments and checks

apache · Jan 29, 2015 · 851f2a2 · 851f2a2
1 parent cdb9895
commit 851f2a2
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 7 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/linalg/distributed/BlockMatrix.scala b/mllib/src/main/scala/org/apache/spark/mllib/linalg/distributed/BlockMatrix.scala
@@ -17,15 +17,15 @@
 
 package org.apache.spark.mllib.linalg.distributed
 
+import scala.collection.mutable.ArrayBuffer
+
 import breeze.linalg.{DenseMatrix => BDM}
 
 import org.apache.spark.{Logging, Partitioner}
 import org.apache.spark.mllib.linalg.{SparseMatrix, DenseMatrix, Matrix}
 import org.apache.spark.rdd.RDD
 import org.apache.spark.storage.StorageLevel
 
-import scala.collection.mutable.ArrayBuffer
-
 /**
  * A grid partitioner, which uses a regular grid to partition coordinates.
  *
@@ -189,18 +189,22 @@ class BlockMatrix(
     val entryRDD = blocks.flatMap { case ((blockRowIndex, blockColIndex), mat) =>
       val rowStart = blockRowIndex.toLong * rowsPerBlock
       val colStart = blockColIndex.toLong * colsPerBlock
-      val entryValues = new ArrayBuffer[MatrixEntry](mat.numRows * mat.numCols)
       mat match {
         case dn: DenseMatrix =>
+          val entryValues = new ArrayBuffer[MatrixEntry](mat.numRows * mat.numCols)
           dn.foreachActive { (i, j, v) =>
             if (v != 0.0) entryValues.append(new MatrixEntry(rowStart + i, colStart + j, v))
           }
+          entryValues
         case sp: SparseMatrix =>
+          val entryValues = new Array[MatrixEntry](sp.values.length)
+          var cnt = 0
           sp.foreachActive { (i, j, v) =>
-            entryValues.append(new MatrixEntry(rowStart + i, colStart + j, v))
+            entryValues(cnt) = new MatrixEntry(rowStart + i, colStart + j, v)
+            cnt += 1
           }
+          entryValues
       }
-      entryValues
     }
     new CoordinateMatrix(entryRDD, numRows(), numCols())
   }

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/linalg/distributed/CoordinateMatrix.scala b/mllib/src/main/scala/org/apache/spark/mllib/linalg/distributed/CoordinateMatrix.scala
@@ -97,8 +97,19 @@ class CoordinateMatrix(
     toIndexedRowMatrix().toRowMatrix()
   }
 
-  /** Converts to BlockMatrix. */
+  /**
+   * Converts to BlockMatrix.
+   * @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 toBlockMatrix(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 numRowBlocks = math.ceil(m.toDouble / rowsPerBlock).toInt

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/linalg/distributed/IndexedRowMatrix.scala b/mllib/src/main/scala/org/apache/spark/mllib/linalg/distributed/IndexedRowMatrix.scala
@@ -77,7 +77,14 @@ class IndexedRowMatrix(
     new RowMatrix(rows.map(_.vector), 0L, nCols)
   }
 
-  /** Converts to BlockMatrix. */
+  /**
+   * Converts to BlockMatrix.
+   * @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 toBlockMatrix(rowsPerBlock: Int, colsPerBlock: Int): BlockMatrix = {
     toCoordinateMatrix().toBlockMatrix(rowsPerBlock, colsPerBlock)
   }

diff --git a/mllib/src/test/scala/org/apache/spark/mllib/linalg/distributed/CoordinateMatrixSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/linalg/distributed/CoordinateMatrixSuite.scala
@@ -106,5 +106,12 @@ class CoordinateMatrixSuite extends FunSuite with MLlibTestSparkContext {
     assert(blockMat.numRows() === m)
     assert(blockMat.numCols() === n)
     assert(blockMat.toBreeze() === mat.toBreeze())
+
+    intercept[IllegalArgumentException] {
+      mat.toBlockMatrix(-1, 2)
+    }
+    intercept[IllegalArgumentException] {
+      mat.toBlockMatrix(2, 0)
+    }
   }
 }
diff --git a/mllib/src/test/scala/org/apache/spark/mllib/linalg/distributed/IndexedRowMatrixSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/linalg/distributed/IndexedRowMatrixSuite.scala
@@ -94,6 +94,13 @@ class IndexedRowMatrixSuite extends FunSuite with MLlibTestSparkContext {
     assert(blockMat.numRows() === m)
     assert(blockMat.numCols() === n)
     assert(blockMat.toBreeze() === idxRowMat.toBreeze())
+
+    intercept[IllegalArgumentException] {
+      idxRowMat.toBlockMatrix(-1, 2)
+    }
+    intercept[IllegalArgumentException] {
+      idxRowMat.toBlockMatrix(2, 0)
+    }
   }
 
   test("multiply a local matrix") {