Search code examples
apache-sparkpysparkapache-spark-sqlsparkr

Getting difference between value and its lag in Spark


I have a SparkR DataFrame as shown below. I want to create a monthdiff column that is the months between dates, grouped by each name. How can I do this?

#Set up data frame
team <- data.frame(name = c("Thomas", "Thomas", "Thomas", "Thomas", "Bill", "Bill", "Bill"),
  dates = c('2017-01-05', '2017-02-23', '2017-03-16', '2017-04-08', '2017-06-08','2017-07-24','2017-09-05'))
#Create Spark DataFrame
team <- createDataFrame(team)
#Convert dates to date type
team <- withColumn(team, 'dates', cast(team$dates, 'date'))

Here's what I've tried so far, all resulting in errors:

team <- agg(groupBy(team, 'name'), monthdiff=c(NA, months_between(team$dates, lag(team$dates))))
team <- agg(groupBy(team, 'name'), monthdiff=months_between(team$dates, lag(team$dates)))
team <- agg(groupBy(team, 'name'), monthdiff=months_between(select(team, 'dates'), lag(select(team, 'dates'))))

Expected output:

name    | dates     | monthdiff
-------------------------------
Thomas  |2017-01-05 |  NA
Thomas  |2017-02-23 |  1
Thomas  |2017-03-16 |  1
Thomas  |2017-04-08 |  1
Bill    |2017-06-08 |  NA
Bill    |2017-07-24 |  1
Bill    |2017-09-05 |  2

Solution

  • Based on this post, I adapted the code for SparkR to get the answer.

    #Create 'lagdates' variable with lag of dates
    window <- orderBy(windowPartitionBy("name"), team$dates)
    team <- withColumn(team, 'lagdates', over(lag(team$dates), window))
    
    #Get months_between dates and lagdates
    team <- withColumn(team, 'monthdiff', round(months_between(team$dates, team$lagdates)))
    
    name  | dates      | lagdates  | monthdiff
    ------------------------------------------
    Bill  | 2017-06-08 |null       | null
    Bill  | 2017-07-24 |2017-06-08 |    2
    Bill  | 2017-09-05 |2017-07-24 |    1
    Thomas| 2017-01-05 |null       | null
    Thomas| 2017-02-23 |2017-01-05 |    2
    Thomas| 2017-03-16 |2017-02-23 |    1
    Thomas| 2017-04-08 |2017-03-16 |    1