pyspark: calculate row difference with variable length partitions

user_384803 0
Viewed 23

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'}
]
pyspark
1 Answers

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')))
     )
user_29561350
Related