title: Pyspark window function with condition
body: <p>Suppose I have a DataFrame of events with time difference between each row, the main rule is that one visit is counted if only the event has been within 5 minutes of the previous or next event:</p>
  
  <pre class="lang-python prettyprint-override"><code>+--------+-------------------+--------+
  |userid  |eventtime          |timeDiff|
  +--------+-------------------+--------+
  |37397e29|2017-06-04 03:00:00|60      |
  |37397e29|2017-06-04 03:01:00|60      |
  |37397e29|2017-06-04 03:02:00|60      |
  |37397e29|2017-06-04 03:03:00|180     |
  |37397e29|2017-06-04 03:06:00|60      |
  |37397e29|2017-06-04 03:07:00|420     |
  |37397e29|2017-06-04 03:14:00|60      |
  |37397e29|2017-06-04 03:15:00|1140    |
  |37397e29|2017-06-04 03:34:00|540     |
  |37397e29|2017-06-04 03:53:00|540     |
  +--------+----------------- -+--------+
  </code></pre>
  
  <p>The challenge is to group by the start_time and end_time of the latest eventtime that has the condition of being within 5 minutes. The output should be like this table:</p>
  
  <pre class="lang-python prettyprint-override"><code>+--------+-------------------+--------------------+-----------+
  |userid  |start_time         |end_time            |events     |
  +--------+-------------------+--------------------+-----------+
  |37397e29|2017-06-04 03:00:00|2017-06-04 03:07:00 |6          |
  |37397e29|2017-06-04 03:14:00|2017-06-04 03:15:00 |2          |
  +--------+-------------------+--------------------+-----------+
  </code></pre>
  
  <p>So far I have used window lag functions and some conditions, however, I do not know where to go from here:</p>
  
  <pre class="lang-python prettyprint-override"><code>%spark.pyspark
  from pyspark.sql import functions as F
  from pyspark.sql import Window as W
  from pyspark.sql.functions import col
  
  windowSpec = W.partitionBy(result_poi["userid"], result_poi["unique_reference_number"]).orderBy(result_poi["eventtime"])
  windowSpecDesc = W.partitionBy(result_poi["userid"], result_poi["unique_reference_number"]).orderBy(result_poi["eventtime"].desc())
  
  # The windows are between the current row and following row. e.g: 3:00pm and 3:03pm 
  nextEventTime = F.lag(col("eventtime"), -1).over(windowSpec)
  
  # The windows are between the current row and following row. e.g: 3:00pm and 3:03pm 
  previousEventTime = F.lag(col("eventtime"), 1).over(windowSpec)
  diffEventTime = nextEventTime - col("eventtime")
  
  nextTimeDiff = F.coalesce((F.unix_timestamp(nextEventTime)
              - F.unix_timestamp('eventtime')), F.lit(0))
  previousTimeDiff = F.coalesce((F.unix_timestamp('eventtime') -F.unix_timestamp(previousEventTime)), F.lit(0))
  
  
  # Check if the next POI is the equal to the current POI and has a time differnce less than 5 minutes.
  validation = F.coalesce(( (nextTimeDiff &lt; 300) | (previousTimeDiff &lt; 300) ), F.lit(False))
  
  # Change True to 1
  visitCheck = F.coalesce((validation == True).cast("int"), F.lit(1))
  
  
  result_poi.withColumn("visit_check", visitCheck).withColumn("nextTimeDiff", nextTimeDiff).select("userid", "eventtime", "nextTimeDiff", "visit_check").orderBy("eventtime")
  </code></pre>
  
  <p>My questions: Is this a viable approach, and if so, how can I "go forward" and look at the maximum eventtime that fulfill the 5 minutes condition. To my knowledge, iterate through values of a Spark SQL Column, is it possible? wouldn't it be too expensive?. Is there another way to achieve this result?</p>
  
  <p><strong>Result of Solution suggested by @Aku</strong>:</p>
  
  <pre class="lang-python prettyprint-override"><code>+--------+--------+---------------------+---------------------+------+
  |userid  |subgroup|start_time           |end_time             |events|
  +--------+--------+--------+------------+---------------------+------+
  |37397e29|0       |2017-06-04 03:00:00.0|2017-06-04 03:06:00.0|5     |
  |37397e29|1       |2017-06-04 03:07:00.0|2017-06-04 03:14:00.0|2     |
  |37397e29|2       |2017-06-04 03:15:00.0|2017-06-04 03:15:00.0|1     |
  |37397e29|3       |2017-06-04 03:34:00.0|2017-06-04 03:43:00.0|2     |
  +------------------------------------+-----------------------+-------+
  </code></pre>
  
  <p>It doesn't give the result expected. 3:07 - 3:14 and  03:34-03:43 are being counted as ranges within 5 minutes, it shouldn't be like that. Also, 3:07 should be the end_time in the first row as it is within 5 minutes of the previous row 3:06.</p>