Scala Memory leak in fo/ yield

Question

Hitting a memory issue with the below sparse matrix multiply routine.

Works fine for very large sparse very sparse matrices, but as the density of the matrices increases it grinds to a halt.

Looks like it leaking memory here.

kk <- (A.colPtrs(jj) to (A.colPtrs(jj + 1)-1)) // fails over here

def sparseMatrixMultiply(A:CSCMatrix[Double], B:CSCMatrix[Double]) = {
      val n = B.cols
      val Bvals = B.activeValuesIterator.toArray
      val Avals = A.activeValuesIterator.toArray

      val entries =
          for {
            j <- (0 to (n-1))/*.par*/
            k <- (B.colPtrs(j) to (B.colPtrs(j + 1)-1))
            jj = B.rowIndices(k) 
            kk <- (A.colPtrs(jj) to (A.colPtrs(jj + 1)-1)) //  get stuck here 
            i = A.rowIndices(kk) 
          } yield {
            println(s"$j $k $jj $kk $i ${Bvals(k)} ${Avals(kk)}")
            (i, j, Bvals(k) * Avals(kk))
          }
        }

      val coo = SparseMatrix.fromCOO(A.rows,B.cols, entries/*.seq*/)

      //converting back to breeze CSCMatrix
      new breeze.linalg.CSCMatrix[Double](coo.values, coo.numRows, coo.numCols, coo.colPtrs, coo.rowIndices)
    }

I have tried splitting it into separate blocks, but iteration quickly runs out of memory at the same step

kk <- (A.colPtrs(jj) to (A.colPtrs(jj + 1)-1))

Is it just allocating memory too fast for GC ? or are all the range instances held onto until the for/yield completes

looking for insight into how the GC manages memory in for/yeild comprehension

Thanks

user_1305121 · Answer

for comprehension is syntax for nested maps and flatMaps. What you did is:

for {
  j <- (0 to (n-1))/*.par*/
  k <- (B.colPtrs(j) to (B.colPtrs(j + 1)-1))
  jj = B.rowIndices(k) 
  kk <- (A.colPtrs(jj) to (A.colPtrs(jj + 1)-1)) //  get stuck here 
  i = A.rowIndices(kk) 
} yield {
  println(s"$j $k $jj $kk $i ${Bvals(k)} ${Avals(kk)}")
  (i, j, Bvals(k) * Avals(kk))
}

which is

(0 to (n-1)).flatMap { j =>
  (B.colPtrs(j) to (B.colPtrs(j + 1)-1)).flatMap { k =>
    val jj = B.rowIndices(k)
    (A.colPtrs(jj) to (A.colPtrs(jj + 1)-1)).map { kk =>
      val i = A.rowIndices(kk)
      println(s"$j $k $jj $kk $i ${Bvals(k)} ${Avals(kk)}")
      (i, j, Bvals(k) * Avals(kk))
    }
  }
}

Each nesting is a closure which:

stores all variables from outer nesting existing when it was created (so everything that appears above it in code)
have to compute a whole sub-result
can be garbage collected in whole only after this subresult is returned and incorporated into the internal builder of a bigger result existing in the outer context

So if you e.g. have 10 items on average introduced by each nesting you are creating a product of 10 x 10 x 10 items through intermediate steps (10 items, each having a closure which creates a 10 items, each creates a closure, that creates the final result, and innermost has to be flattened, and then middle nesting has to be flattened, and the the outermost result has to be flattened). Though your issues suggest that the amount of intermediate objects is way above the order of 1000s.

This is a lot of intermediate results AND this is not a memory leak. This is how it is supposed to work. All that memory is collected once you exit the for comprehension. Ranges are light on their own, but when you use .map and .flatMap you convert them into an IndexedSeq which will most likely be some eager, memoizing data structure that would compute the final result at once without allowing JVM any abilities to forget anything (since interface only guarantees IndexedSeq it is difficult to assume which one, it might be List or Vector but you have no guarantees).

You might try converting each Range into LazyList (2.13) or Stream (2.12 and before) and before passing into fromCOO converting this lazy data structure into Iterator or Iterable. This way for comprehension would build a structure which allocates and remembers only as much as needed to compute the next element and immediately forgets everything that is no longer needed. Something like:

val entries = for {
  j <- (0 to (n-1)).to(LazyList)
  k <- (B.colPtrs(j) to (B.colPtrs(j + 1)-1)).to(LazyList)
  jj = B.rowIndices(k) 
  kk <- (A.colPtrs(jj) to (A.colPtrs(jj + 1)-1)).to(LazyList)
  i = A.rowIndices(kk) 
} yield {
  println(s"$j $k $jj $kk $i ${Bvals(k)} ${Avals(kk)}")
  (i, j, Bvals(k) * Avals(kk))
}

SparseMatrix.fromCOO(A.rows,B.cols, entries.toIterable)

Scala Memory leak in fo/ yield

1 Answers