Question: Suppose I have the following dataframe in pyspark:     object time has_changed     A 1 0   A 2 1   A 4 0   A 7 1   B 2 1   B 5 0     What I want is to add a new column that, for each row, keeps track of the time difference with respect to the last value change for the current object (or first element of the corresponding partition if no value changes exists). For the table I've posted above, the result would be the following:     object time has_changed time_alive     A 1 0 0   A 2 1 1   A 4 0 2   A 7 1 5   B 2 1 0   B 5 0 3     That is, within each partition by the "object" column, sorted by the "time" column, each value of the corresponding row is calculated as the difference between the time of that row and the previous time at which there is a 1 in the "has_changed" column (if a 1 is not found, the window will scroll to the first element of the partition). What I would like to implement would be something like the following (pseudo-code): from pyspark.sql.window import Window as w from pyspark.sql import SparkSession spark = SparkSession.builder.getOrCreate() # Define the data data = [("A", 1, 0), ("A", 2, 1), ("A", 4, 0), ("A", 7, 1), ("B", 2, 1), ("B", 5, 0)] # Define the schema schema = ["object", "time", "has_changed"] # Create the DataFrame df = spark.createDataFrame(data, schema) # Window function (pseudo-code, this won't work) window = ( w.partitionBy("object") .orderBy("time") .rowsBetween(f.when(f.col("has_changed") == 1), w.currentRow) ) df.withColumn("time_alive", f.col("time") - f.lag("time", 1).over(window)) 

Answer: Step by step solution Create a window specification W = Window.partitionBy('object').orderBy('time')  Mask the values in time column where has_changed is 0 masked = F.when(F.col('has_changed') == 1, F.col('time')) df = df.withColumn('masked', masked) # +------+----+-----------+------+ # |object|time|has_changed|masked| # +------+----+-----------+------+ # | A| 1| 0| NULL| # | A| 2| 1| 2| # | A| 4| 0| NULL| # | A| 7| 1| 7| # | B| 2| 1| 2| # | B| 5| 0| NULL| # +------+----+-----------+------+  Calculate the first value in time per group df = df.withColumn('first', F.first('time').over(W)) # +------+----+-----------+------+-----+ # |object|time|has_changed|masked|first| # +------+----+-----------+------+-----+ # | A| 1| 0| NULL| 1| # | A| 2| 1| 2| 1| # | A| 4| 0| NULL| 1| # | A| 7| 1| 7| 1| # | B| 2| 1| 2| 2| # | B| 5| 0| NULL| 2| # +------+----+-----------+------+-----+  Forward fill and shift the last valid value in masked time column over the window last_changed = F.lag(F.last('masked', ignorenulls=True).over(W)).over(W) df = df.withColumn('last_changed', last_changed) # +------+----+-----------+------+-----+------------+ # |object|time|has_changed|masked|first|last_changed| # +------+----+-----------+------+-----+------------+ # | A| 1| 0| NULL| 1| NULL| # | A| 2| 1| 2| 1| NULL| # | A| 4| 0| NULL| 1| 2| # | A| 7| 1| 7| 1| 2| # | B| 2| 1| 2| 2| NULL| # | B| 5| 0| NULL| 2| 2| # +------+----+-----------+------+-----+------------+  Fill the nulls in last_changed with the first value in group last_changed = F.when(last_changed.isNull(), first).otherwise(last_changed) df = df.withColumn('last_changed', last_changed) # +------+----+-----------+------+-----+------------+ # |object|time|has_changed|masked|first|last_changed| # +------+----+-----------+------+-----+------------+ # | A| 1| 0| NULL| 1| 1| # | A| 2| 1| 2| 1| 1| # | A| 4| 0| NULL| 1| 2| # | A| 7| 1| 7| 1| 2| # | B| 2| 1| 2| 2| 2| # | B| 5| 0| NULL| 2| 2| # +------+----+-----------+------+-----+------------+  Subtract time column from last_changed to calculate time_alive df = df.withColumn('time_alive', F.col('time') - last_changed) # +------+----+-----------+------+-----+------------+----------+ # |object|time|has_changed|masked|first|last_changed|time_alive| # +------+----+-----------+------+-----+------------+----------+ # | A| 1| 0| NULL| 1| 1| 0| # | A| 2| 1| 2| 1| 1| 1| # | A| 4| 0| NULL| 1| 2| 2| # | A| 7| 1| 7| 1| 2| 5| # | B| 2| 1| 2| 2| 2| 0| # | B| 5| 0| NULL| 2| 2| 3| # +------+----+-----------+------+-----+------------+----------+