Skip to content

Commit 41d0645

Browse files
srowenkai-chi
authored andcommitted
[SPARK-26228][MLLIB] OOM issue encountered when computing Gramian matrix
Avoid memory problems in closure cleaning when handling large Gramians (>= 16K rows/cols) by using null as zeroValue Existing tests. Note that it's hard to test the case that triggers this issue as it would require a large amount of memory and run a while. I confirmed locally that a 16K x 16K Gramian failed with tons of driver memory before, and didn't fail upfront after this change. Closes apache#23600 from srowen/SPARK-26228. Authored-by: Sean Owen <sean.owen@databricks.com> Signed-off-by: Sean Owen <sean.owen@databricks.com> (cherry picked from commit 6dcad38) Signed-off-by: Sean Owen <sean.owen@databricks.com>
1 parent 0e1d0bb commit 41d0645

1 file changed

Lines changed: 17 additions & 3 deletions

File tree

  • mllib/src/main/scala/org/apache/spark/mllib/linalg/distributed

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

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -119,11 +119,25 @@ class RowMatrix @Since("1.0.0") (
119119
val nt = if (n % 2 == 0) ((n / 2) * (n + 1)) else (n * ((n + 1) / 2))
120120

121121
// Compute the upper triangular part of the gram matrix.
122-
val GU = rows.treeAggregate(new BDV[Double](nt))(
123-
seqOp = (U, v) => {
122+
val GU = rows.treeAggregate(null.asInstanceOf[BDV[Double]])(
123+
seqOp = (maybeU, v) => {
124+
val U =
125+
if (maybeU == null) {
126+
new BDV[Double](nt)
127+
} else {
128+
maybeU
129+
}
124130
BLAS.spr(1.0, v, U.data)
125131
U
126-
}, combOp = (U1, U2) => U1 += U2)
132+
}, combOp = (U1, U2) =>
133+
if (U1 == null) {
134+
U2
135+
} else if (U2 == null) {
136+
U1
137+
} else {
138+
U1 += U2
139+
}
140+
)
127141

128142
RowMatrix.triuToFull(n, GU.data)
129143
}

0 commit comments

Comments
 (0)