title: How to assign ranks to records in a spark dataframe based on some conditions?
body: <p>Given a dataframe :</p>
  
  <pre><code>+-------+-------+
  |   A   |   B   |
  +-------+-------+
  |      a|      1|
  +-------+-------+
  |      b|      2|
  +-------+-------+
  |      c|      5|
  +-------+-------+
  |      d|      7|
  +-------+-------+
  |      e|     11|
  +-------+-------+    
  </code></pre>
  
  <p>I want to assign ranks to records based on conditions :</p>
  
  <ol>
  <li>Start rank with 1</li>
  <li>Assign rank = rank of previous record if ( B of current record - B of previous record ) is &lt;= 2</li>
  <li>Increment rank when ( B of current record - B of previous record ) is > 2</li>
  </ol>
  
  <p>So I want result to be like this : </p>
  
  <pre><code>+-------+-------+------+
  |   A   |   B   | rank |
  +-------+-------+------+
  |      a|      1|     1|
  +-------+-------+------+
  |      b|      2|     1|
  +-------+-------+------+
  |      c|      5|     2|
  +-------+-------+------+
  |      d|      7|     2|
  +-------+-------+------+
  |      e|     11|     3|
  +-------+-------+------+
  </code></pre>
  
  <ul>
  <li>Inbuilt functions in spark like rowNumber, rank, dense_rank don't
  provide any functionality to achieve this.</li>
  <li>I tried doing it by using a global variable rank and fetching
  previous record values using lag function but it does not give
  consistent results due to distributed processing in spark unlike in sql.</li>
  <li><p>One more method I tried was passing lag values of records to a UDF while generating a new column and applying conditions in UDF. But the problem I am facing is I can get lag values for columns A as well as B but not for column rank.
  This gives error as it cannot resolve column name rank :</p>
  
  <p>HiveContext.sql("SELECT df.*,LAG(df.rank, 1) OVER (ORDER BY B , 0) AS rank_lag, udfGetVisitNo(B,rank_lag) as rank  FROM df") </p></li>
  <li><p>I cannot get lag value of a column which I am currently adding.</p></li>
  <li><p>Also I dont want methods which require using df.collect() as this dataframe is quite large in size and collecting it on a single working node results in memory errors.</p></li>
  </ul>
  
  <p>Any other method by which I can achieve the same?
  I would like to know a solution having time complexity O(n) , n being the no of records.</p>