title: PySpark: How to add columns for the next and previous event in a sorted database?
body: <p>Let's say I've got dataset like this one:</p>
  <pre><code>  +----------+------------+----------+ -----------+
    |userId     |event      |timestamp | session_id |
    +----------+------------+----------+ -----------+
    |U1        |A           |1         | 11         |
    |U1        |B           |2         | 11         |
    |U1        |C           |5         | 11         |
    |U1        |C           |7         | 12         |
    |U2        |A           |8         | 13         |
    |U2        |D           |20        | 13         |
    |U3        |B           |23        | 14         |
    +----------+------------+----------+------------+
  
  </code></pre>
  <p>The dataset is sorted based on the session_id and the timestamp. I would like to add a column called &quot;next_event&quot; where it should be placed the next event in the session (so in this database the event in the row below) if it exists. If there is no next event in the session, it should put &quot;NA&quot; or something similar.</p>
  <p>Similarly, I would like to add a column called &quot;prev_event&quot; where it should be placed the previous event in the session (so in this database the event in the row above) if it exists. If there is no previous event in the session, it should put &quot;NA&quot; or something similar.</p>
  <p>This is optional, but it would be very helpful if I could also add a column called &quot;uniq_id&quot; from the beginning of the dataset until the end that would be numbered from 0 to n-1 or from 1 to n. And another column &quot;next_id&quot; similar to &quot;next_event&quot; but that instead of the next_event specifies the next uniq_id. This isn't necessary, but I think it might help me later on.</p>
  <pre><code>   +----------+------------+----------+ -----------+------------+------------+---------+---------+
    |userId    |event      |timestamp | session_id | next_event  | prev_event | uniq_id | next_id |
    +----------+------------+----------+ -----------+-------------+------------+---------+---------+
    |U1        |A           |1         | 11         |   B         |   NA       |  1      | 2       |
    |U1        |B           |2         | 11         |   C         |   A        |  2      | 3       |
    |U1        |C           |5         | 11         |   NA        |   B        |  3      | NA      |
    |U1        |C           |7         | 12         |   NA        |   NA       |  4      | NA      |
    |U2        |A           |8         | 13         |   D         |   NA       |  5      | 6       |
    |U2        |D           |20        | 13         |   NA        |   A        |  6      | NA      |
    |U3        |B           |23        | 14         |   NA        |   NA       |  7      | NA      |
    +----------+------------+----------+------------+-------------+------------+---------+---------+
  
  </code></pre>
  <p>I would like this to be done in PySpark.
  Any sort of help or guidance is very appreciated! Thanks</p>
  <p>Edit: I think I could use lag and lead with a Window for this, but I'm not exactly sure how</p>