pyspark: calculate row difference with variable length partitions - pyspark

Question

I have a time series dataset where each row is a reading of a certain property. A number of properties are read at the same time in a batch and are tagged with a common time stamp. The number of properties read from time to time can vary (i.e. these batches are of variable length). I need to label each such batch with time delta from the previous batch. I know how to do this for fix-sized batches, but can't figure out how to do this in this particular case.
I need to perform this operation on tens of millions of rows, so the solution has to be highly efficient.
Here's sample data:
from pyspark.sql import SparkSession
spark = SparkSession.builder.appName('sparkdf').getOrCreate()
data_in = [
{'time': 12, 'batch_id': 1, 'name': 'a'}
, {'time': 12, 'batch_id': 1, 'name': 'c'}
, {'time': 12, 'batch_id': 1, 'name': 'e'}
, {'time': 12, 'batch_id': 1, 'name': 'd'}
, {'time': 12, 'batch_id': 1, 'name': 'e'}
, {'time': 14, 'batch_id': 2, 'name': 'a'}
, {'time': 14, 'batch_id': 2, 'name': 'b'}
, {'time': 14, 'batch_id': 2, 'name': 'c'}
, {'time': 19, 'batch_id': 3, 'name': 'b'}
, {'time': 19, 'batch_id': 3, 'name': 'c'}
, {'time': 19, 'batch_id': 3, 'name': 'e'}
, {'time': 19, 'batch_id': 3, 'name': 'f'}
, {'time': 19, 'batch_id': 3, 'name': 'g'}
]
# creating a dataframe
dataframe_in = spark.createDataFrame(data)
# show data frame
display(dataframe_in.select('time', 'sample_id', 'name'))
here's what I'm looking for in the output:
data_out = [
{'time': 12, 'time_delta':None, 'batch_id': 1, 'name': 'a'}
, {'time': 12, 'time_delta':None, 'batch_id': 1, 'name': 'c'}
, {'time': 12, 'time_delta':None, 'batch_id': 1, 'name': 'e'}
, {'time': 12, 'time_delta':None, 'batch_id': 1, 'name': 'd'}
, {'time': 12, 'time_delta':None, 'batch_id': 1, 'name': 'e'}
, {'time': 14, 'time_delta':2, 'batch_id': 2, 'name': 'a'}
, {'time': 14, 'time_delta':2, 'batch_id': 2, 'name': 'b'}
, {'time': 14, 'time_delta':2, 'batch_id': 2, 'name': 'c'}
, {'time': 19, 'time_delta':5, 'batch_id': 3, 'name': 'b'}
, {'time': 19, 'time_delta':5, 'batch_id': 3, 'name': 'c'}
, {'time': 19, 'time_delta':5, 'batch_id': 3, 'name': 'e'}
, {'time': 19, 'time_delta':5, 'batch_id': 3, 'name': 'f'}
, {'time': 19, 'time_delta':5, 'batch_id': 3, 'name': 'g'}
]

Answer 1

You can use first and lag to get the previous group's value.
df = (df.withColumn('lag', F.lag('time').over(Window.orderBy('batch_id')))
.withColumn('time_delta', F.col('time') -
F.first('lag').over(Window.partitionBy('batch_id').orderBy('lag')))
)
Alternatively, you can use min and this might be slightly more efficient because I can omit the orderBy but this will return 0 for batch 1.
df = (df.withColumn('lag', F.lag('time').over(Window.orderBy('batch_id')))
.withColumn('time_delta', F.col('time') -
F.min('lag').over(Window.partitionBy('batch_id')))
)