pyspark: resampling pyspark dataframe on date and time

Viewed 1157

How to resample pyspark dataframe, like in pandas we have pd.grouper, and pd.resample and I can resample on h,2h, 3h, week. I have the following sample pyspark dataframe, how I can aggregate it at on column ind and date and every h/2h/3h

from pyspark import SparkContext
from pyspark.sql import SQLContext

sc = SparkContext.getOrCreate()
sqlContext = SQLContext(sc)

a = sqlContext.createDataFrame([["Anand", "2020-02-01 16:00:00", 12, "ba"],
                         ["Anand", "2020-02-01 16:05:00", 7, "ba" ]
                        ["Anand", "2020-02-02 19:10:00", 14,"sa"], 
                        ["Carl", "2020-02-01 16:00:00", 16,"da"], 
                        ["Carl", "2020-02-02 16:02:00", 12,"ga"],
                        ["Carl", "2020-02-02 17:10:00", 1,"ga"],
                        ["Eric", "2020-02-01 16:o0:00", 24, "sa"]], ['ind',"date","sal","imp"])
a.show()

|  ind|               date|sal|imp|
+-----+-------------------+---+---+
|Anand|2020-02-01 16:00:00| 12| ba|
|Anand|2020-02-01 16:05:00|  7| sa|
|Anand|2020-02-02 19:10:00| 14| sa|
| Carl|2020-02-01 16:00:00| 16| da|
| Carl|2020-02-01 16:02:00| 12| ga|
| Carl|2020-02-02 17:10:00|  1| ga|
| Eric|2020-02-01 16:00:00| 24| sa|

So when aggregation on column ind and resampling on date(everyhour) and mean of sale Desired output may look like

|  ind|               date|sal|
+-----+-------------------+---+
|Anand|2020-02-01 16:00:00|  9|
|Anand|2020-02-02 19:00:00| 14|
| Carl|2020-02-01 16:00:00|  9|
| Carl|2020-02-02 17:00:00|  1|
| Eric|2020-02-01 16:00:00| 24|
3 Answers

You can do exactly what you have already described in your question: group by ind and date. With the help of date_trunc we can round the date column to hours before grouping:

from pyspark.sql import functions as F
a.groupBy('ind', F.date_trunc('hour', F.col('date')).alias('date'))\
   .agg(F.mean('sal')) \
   .orderBy('ind', 'date') \
   .show()

prints

+-----+-------------------+--------+
|  ind|               date|avg(sal)|
+-----+-------------------+--------+
|Anand|2020-02-01 16:00:00|     9.5|
|Anand|2020-02-02 19:00:00|    14.0|
| Carl|2020-02-01 16:00:00|    14.0|
| Carl|2020-02-02 17:00:00|     1.0|
| Eric|2020-02-01 16:00:00|    24.0|
+-----+-------------------+--------+

One possible approach is using 2 windows 1 to determine if time difference over a partition of ind and date is within 1 hr , and second to get mean using the above window and the time_diff calculated (Note: For Anand (12+7)/2 = 9.5 as oppsed to 9 in the expected output):

one_hrs= 1*60*60
w = Window.partitionBy("ind",F.to_date("date"))
w1 = Window.partitionBy("ind",F.to_date("date"),"time_diff")

(df.withColumn("date",F.to_timestamp("date"))
   .withColumn("first_date",F.first("date").over(w))
   .withColumn("time_diff",((F.unix_timestamp("date")-F.unix_timestamp("first_date"))
    <=one_hrs).cast("Integer"))
 .withColumn("sal",F.mean("sal").over(w1)).dropDuplicates(["ind","sal","time_diff"])
 .drop("first_date","time_diff").orderBy("ind")).show()

+-----+-------------------+----+---+
|  ind|               date| sal|imp|
+-----+-------------------+----+---+
|Anand|2020-02-02 19:10:00|14.0| sa|
|Anand|2020-02-01 16:00:00| 9.5| ba|
| Carl|2020-02-01 16:00:00|14.0| da|
| Carl|2020-02-02 17:10:00| 1.0| ga|
| Eric|2020-02-01 16:00:00|24.0| sa|
+-----+-------------------+----+---+

Seeing that you have the date as string, one crude way of doing it would be to split and aggregate.

import pyspark.sql.functions as F
a = sqlContext.createDataFrame([["Anand", "2020-02-01 16:00:00", 12, "ba"],
                         ["Anand", "2020-02-01 16:05:00", 7, "ba"],
                        ["Anand", "2020-02-02 19:10:00", 14,"sa"], 
                        ["Carl", "2020-02-01 16:00:00", 16,"da"], 
                        ["Carl", "2020-02-02 16:02:00", 12,"ga"],
                        ["Carl", "2020-02-02 17:10:00", 1,"ga"],
                        ["Eric", "2020-02-01 16:o0:00", 24, "sa"]], ['ind',"date","sal","imp"])
a_spli = a.withColumn("hour",F.split(F.col('date'),':')[0])

test_res = a_spli.groupby('ind','hour').agg(F.mean('sal'))

sparkts is a cool library to handle time related tasks: https://github.com/sryza/spark-timeseries. check it out.

test_res.show()
+-----+-------------+--------+
|  ind|         hour|avg(sal)|
+-----+-------------+--------+
|Anand|2020-02-01 16|     9.5|
|Anand|2020-02-02 19|    14.0|
| Carl|2020-02-01 16|    16.0|
| Carl|2020-02-02 16|    12.0|
| Carl|2020-02-02 17|     1.0|
| Eric|2020-02-01 16|    24.0|
+-----+-------------+--------+
Related