Pyspark - How to aggregate over '4 hours' windows in groups

The idea is to partition the data by email, sort within each partition by date and time and then map each partition to the desired output. This approach will work if the data for each partition (=data for one email adress) fits into the memory of one Spark executor.

The acutal Spark logic follows the steps

1. Create a new column which contains the timestamps
2. Partition the data by email, so that all rows with the same email are part of the same partition. Note that there might be data from more than one email in one partition.
3. Sort each partition by email and timestamp.
4. Process each partition. If necessary produce multiple outputs per partition as required

from pyspark.sql.functions import *
from pyspark.sql import SparkSession
from pyspark.sql.types import Row
from datetime import datetime, timedelta

spark = SparkSession.builder.appName("test").getOrCreate()
df = spark.read.option("header", "true").csv(<path>) #or any other data source
df = df.withColumn("date_time", to_timestamp(concat(col("Date_of_purchase"), lit(" "), col("time_of_purchase")), "MM/dd/yy hh:mm aa")) \
    .drop("Date_of_purchase", "time_of_purchase") \
    .repartition(col("email")) \
    .sortWithinPartitions(col("email"), col("date_time"))

def process_partition(df_chunk):
    row_list = list(df_chunk)
    if len(row_list) == 0:
        return
    email = row_list[0]['email']
    start = row_list[0]['date_time']
    end = start + timedelta(hours=4)
    count = 0
    for row in row_list:
        if email == row['email'] and end > row['date_time']:
            count = count +1
        else:
            yield Row(email, start, end, count)
            email = row['email']
            start = row['date_time']
            end = start + timedelta(hours=4)
            count = 1
    yield Row(email, start, end, count)

result = df.rdd.mapPartitions(process_partition).toDF(["email", "from", "to", "count"])
result.show()

Output:

+-------------+-------------------+-------------------+-----+
|        email|               from|                 to|count|
+-------------+-------------------+-------------------+-----+
|def@gmail.com|2018-11-10 12:17:00|2018-11-10 16:17:00|    2|
|def@gmail.com|2018-11-10 20:16:00|2018-11-11 00:16:00|    3|
|abc@gmail.com|2018-11-10 12:10:00|2018-11-10 16:10:00|    3|
|abc@gmail.com|2018-11-11 06:16:00|2018-11-11 10:16:00|    2|
+-------------+-------------------+-------------------+-----+

To change the length of the period, the timedeltas can be set to any value.