title: Spark out of memory with a large number of window functions (lag, lead)
body: <p>I need to calculate additional features from a dataset using multiple lead's and lag's. The high number of lead's and lag's causes a out-of-memory error.</p>
  
  <p>Data frame:</p>
  
  <pre><code>|----------+----------------+---------+---------+-----+---------|
  | DeviceID | Timestamp      | Sensor1 | Sensor2 | ... | Sensor9 |
  |----------+----------------+---------+---------+-----+---------|
  |          |                |         |         |     |         |
  | Long     | Unix timestamp | Double  | Double  |     | Double  |
  |          |                |         |         |     |         |
  |----------+----------------+---------+---------+-----+---------|
  </code></pre>
  
  <p>Window definition:</p>
  
  <pre><code>// Each window contains about 600 rows
  val w = Window.partitionBy("DeviceID").orderBy("Timestamp") 
  </code></pre>
  
  <p>Compute extra features:</p>
  
  <pre><code>var res = df
  val sensors = (1 to 9).map(i =&gt; s"Sensor$i")
  
  for (i &lt;- 1 to 5) {
    for (s &lt;- sensors) {
      res = res.withColumn(lag(s, i).over(w))
           .withColumn(lead(s, i)).over(w)
    }
  
    // Compute features from all the lag's and lead's
    [...]
  }
  </code></pre>
  
  <p>System info:</p>
  
  <pre><code>RAM: 16G
  JVM heap: 11G
  </code></pre>
  
  <p>The code gives correct results with small datasets, but gives an out-of-memory error with 10GB of input data. 
  I think the culprit is the high number of window functions because the DAG shows a very long sequence of</p>
  
  <pre><code>Window -&gt; WholeStageCodeGen -&gt; Window -&gt; WholeStageCodeGen ...
  </code></pre>
  
  <p>Is there anyway to calculate the same features in a more efficient way? 
  For example, is it possible to get lag(Sensor1, 1), lag(Sensor2, 1), ..., lag(Sensor9, 1) without calling lag(..., 1) nine times? </p>
  
  <p>If the answer to the previous question is no, then how can I avoid out-of-memory? I have already tried increasing the number of partitions.</p>