Spark-sql Window Functions On Dataframe - Finding First Timestamp In A Group

I have below dataframe (say UserData). uid region timestamp a 1 1 a 1 2 a 1 3 a 1 4 a 2 5 a 2 6 a 2 7 a 3 8 a 4 9 a 4 10 a 4 11 a 4

Solution 1:

Window functions are indeed useful although your approach can work only if you assume that user visits given region only once. Also window definition you use is incorrect - multiple calls to partitionBy simply return new objects with different window definitions. If you want to partition by multiple columns you should pass them in a single call (.partitionBy("region", "uid")).

Lets start with marking continuous visits in each region:

import org.apache.spark.sql.functions.{lag, sum, not}
import org.apache.spark.sql.expressions.Window 

valw= Window.partitionBy($"uid").orderBy($"timestamp")

valchange= (not(lag($"region", 1).over(w) <=> $"region")).cast("int")
valind= sum(change).over(w)

valdfWithInd= df.withColumn("ind", ind)

Next you we simply aggregate over the groups and find leads:

import org.apache.spark.sql.functions.{lead, coalesce}

valregionTimeEnd= coalesce(lead($"timestamp", 1).over(w), $"max_")

valresult= dfWithInd
  .groupBy($"uid", $"region", $"ind")
  .agg(min($"timestamp").alias("timestamp"), max($"timestamp").alias("max_"))
  .drop("ind")
  .withColumn("regionTimeEnd", regionTimeEnd)
  .withColumnRenamed("timestamp", "regionTimeStart")
  .drop("max_")

result.show

// +---+------+---------------+-------------+// |uid|region|regionTimeStart|regionTimeEnd|// +---+------+---------------+-------------+// |  a|     1|              1|            5|// |  a|     2|              5|            8|// |  a|     3|              8|            9|// |  a|     4|              9|           13|// |  a|     1|             13|           15|// |  a|     3|             15|           17|// |  a|     5|             17|           20|// +---+------+---------------+-------------+