diff --git a/mllib/src/main/scala/org/apache/spark/mllib/linalg/Matrices.scala b/mllib/src/main/scala/org/apache/spark/mllib/linalg/Matrices.scala index c02ba426fcc3a6580a1fb88877874020dc38b9ba..cfed9ad0730bb2537c4cdb684a2cedb4334950d5 100644 --- a/mllib/src/main/scala/org/apache/spark/mllib/linalg/Matrices.scala +++ b/mllib/src/main/scala/org/apache/spark/mllib/linalg/Matrices.scala @@ -296,6 +296,8 @@ class DenseMatrix @Since("1.3.0") ( override def apply(i: Int, j: Int): Double = values(index(i, j)) private[mllib] def index(i: Int, j: Int): Int = { + require(i < numRows && i >=0, s"Expected 0 <= i < $numRows, got $i") + require(j < numCols && j >=0, s"Expected 0 <= j < $numCols, got $j") if (!isTransposed) i + numRows * j else j + numCols * i } @@ -570,6 +572,8 @@ class SparseMatrix @Since("1.3.0") ( } private[mllib] def index(i: Int, j: Int): Int = { + require(i < numRows && i >=0, s"Expected 0 <= i < $numRows, got $i") + require(j < numCols && j >=0, s"Expected 0 <= j < $numCols, got $j") if (!isTransposed) { Arrays.binarySearch(rowIndices, colPtrs(j), colPtrs(j + 1), i) } else { diff --git a/mllib/src/test/scala/org/apache/spark/mllib/linalg/MatricesSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/linalg/MatricesSuite.scala index bfd6d5495f5e0b138d8825b631f95619966655f8..b0071c9a028a2aa28feae1e27442ecc4c0a37492 100644 --- a/mllib/src/test/scala/org/apache/spark/mllib/linalg/MatricesSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/mllib/linalg/MatricesSuite.scala @@ -74,6 +74,17 @@ class MatricesSuite extends SparkFunSuite { } } + test("index in matrices incorrect input") { + val sm = Matrices.sparse(3, 2, Array(0, 2, 3), Array(1, 2, 1), Array(0.0, 1.0, 2.0)) + val dm = Matrices.dense(3, 2, Array(0.0, 2.3, 1.4, 3.2, 1.0, 9.1)) + Array(sm, dm).foreach { mat => + intercept[IllegalArgumentException] { mat.index(4, 1) } + intercept[IllegalArgumentException] { mat.index(1, 4) } + intercept[IllegalArgumentException] { mat.index(-1, 2) } + intercept[IllegalArgumentException] { mat.index(1, -2) } + } + } + test("equals") { val dm1 = Matrices.dense(2, 2, Array(0.0, 1.0, 2.0, 3.0)) assert(dm1 === dm1)