title: Spark: break partition iterator for better memory management?
body: <p>I'm trying to develop a heavy mathematical calculations in Spark, both in term of time and memory (up to <code>O(n^2)</code> for both of them). I've found that the partition holding an <code>Iterator</code> is not really adequate for big calculus since it forces to instantiate (though lazily since it's an <code>Iterator</code>) one object per line. Indeed in a most simple scenario, one would hold a vector per line for instance. But it's both harmful for memory, as we know the JVM overhead for objects and all the pressure that is put on the GC, and for speed, as I may really improve performances improving my linear algebra operations up to BLAS level-3 (matrix by matrix instead of matrix by vector which I'm stuck with in this paradigm). In a very schematic here's what I want to achieve:</p>
  
  <pre><code>while (???) { // loop over some condition, doesn't really matter what
      val matrix = ??? // an instance of a matrix
      val broadMatrix = sparkContext.broadcast(matrix)
      // rdd is an instance of RDD[Vector] that is already cached
      rdd.mapPartition {
          iter =&gt;
              val matrixValue = broadMatrix.value()
              iter.map (vector =&gt; matrixValue * vec)
      }
      // a bunch of other things relying on that result
  }
  </code></pre>
  
  <p>Here are my thoughts:</p>
  
  <ol>
  <li><p>as my <code>rdd</code> in the code above is cached, then having an <code>Iterator</code> is useless, isn't it? Since the only advantage of it is not to hold in memory all the lines at the same time: but here it's been computed and cached so all the lines are held in memory... Yes of course one could argue that Spark's might have an intelligent cache that serializes and compress data (which I doubt when the storage level is <code>MEMORY_ONLY</code> though...).</p></li>
  <li><p>if 1. is true, then the only thing it produces is a huge memory overhead, as I have as many JVM objects as there are rows in my <code>rdd</code> but I could lower it down to a single JVM object per partition. I could even lower it down to a single object per <code>Executor</code> having a scala <code>object</code> that would act as shared memory for all the partitions living on the same executor (this I fear might be hard to handle though as I want to keep Spark's resilience, hence if a partition should be remove for any reason and re-appear on another executor I don't want to handle it by myself but let Spark move all the related objects by itself...).</p></li>
  </ol>
  
  <p>My idea hence would be to transform this <code>rdd</code> of <code>vector</code> into one containing matrices, something like:</p>
  
  <pre><code>while (???) { // loop over some condition, doesn't really matter what
      val matrix = ??? // an instance of a matrix
      val broadMatrix = sparkContext.broadcast(matrix)
      // rdd is an instance of RDD[Vector] that is already cached
      rdd.mapPartition {
          iter =&gt;
              val matrixValue = broadMatrix.value()
              // iter actually contains one single element which is the matrix containing all vectors stacked
              // here we have a BLAS-3 operation
              iter.map (matrix =&gt; matrixValue * matrix)
      }
      // a bunch of other things relying on that result
  }
  </code></pre>
  
  <p>Anyone already faced this dilemna? Have you experienced advance usage of memory management as this one?</p>