Creating a Binary Indicator Column Based on Date Differences in PySpark DataFrame

I am working with a PySpark DataFrame that contains columns 'ID', 'date', and 'bool'. 'bool' is an indicator.

 #+-------+-------------+----+
 #|   ID  |     date    |bool|
 #+-------+-------------+----+
 #|1092829|  2019-10-29 |  0 |
 #|3091902|  2019-12-14 |  1 |
 #|3091902|  2020-07-13 |  0 |
 #|1092829|  2020-07-15 |  1 |
 #|1092829|  2020-07-17 |  1 |
 #|1092829|  2020-08-18 |  1 |
 #|  ...  |     ...     | ...|
 #+-------+-------------+----+

To simplify, let's imagine that I partition my dataframe over the variable ID and, to be even clearer, that I only consider ID 1092829. Note that the DataFrame is ordered in ascending order based on the dates.

 #+-------+-------------+----+
 #|   ID  |     date    |bool|
 #+-------+-------------+----+
 #|1092829|  2019-10-29 |  0 |
 #|1092829|  2019-12-14 |  1 |
 #|1092829|  2020-07-15 |  1 |
 #|1092829|  2020-07-17 |  1 |
 #|1092829|  2020-07-19 |  0 |
 #|1092829|  2020-08-15 |  1 |
 #|1092829|  2020-09-10 |  0 |
 #|1092829|  2020-09-15 |  0 |
 #|1092829|  2020-09-20 |  1 | 
 #|  ...  |     ...     | ...|
 #+-------+-------------+----+

I want to create a column 'D' that is equal to 1 for a line X if the variable bool of a line Y (which is below line X) is equal to 1 and if the date difference between line X and line Y is less than 2 weeks, and 0 otherwhise.

This would yield the following dataframe

 #+-------+-------------+----+---+
 #|   ID  |     date    |bool| D |
 #+-------+-------------+----+---+
 #|1092829|  2019-10-29 |  0 | 0 |
 #|1092829|  2019-12-14 |  1 | 0 |
 #|1092829|  2020-07-15 |  1 | 1 |
 #|1092829|  2020-07-17 |  1 | 0 |
 #|1092829|  2020-07-19 |  0 | 0 |
 #|1092829|  2020-08-15 |  1 | 0 |
 #|1092829|  2020-09-10 |  0 | 1 |
 #|1092829|  2020-09-15 |  0 | 0 |
 #|1092829|  2020-09-20 |  1 | 0 |
 #|  ...  |     ...     | ...|...|
 #+-------+-------------+----+---+

I know how to do it for a fixed number of rows and with the use of pyspark lag function but it's far from optimal as there might be 10 lines of observations in less than a 2 week span in the dataframe.

An example with two lags :

import pyspark.sql.functions as F

df = df.withColumn("D", F.when((F.datediff(df.date, F.lag(df.date,-1))<14) & (F.lag(df.bool,-1)==1)\
| (F.datediff(df.date, F.lag(df.date,-2))<14) & (F.lag(df.bool,-2)==1),1).otherwise(0))

How could I do it for any number of rows in a computationally efficient way?

I think this logic is possible. Count bool = 1 between 2 weeks range and if it is larget than 1, it is the target.

from pyspark.sql import functions as f
from pyspark.sql import Window

w = Window.partitionBy('ID').orderBy('timestamp').rangeBetween(0, 14 * 86400)
​
df.withColumn('timestamp', f.unix_timestamp('date', 'y-M-d')) \
  .withColumn('bool_in_2weeks', f.count(f.when(f.col('bool') == f.lit(1), True)).over(w)) \
  .withColumn('D', f.expr('bool = 1 and bool_in_2weeks > 1').cast('int')) \
  .show(truncate=False)
+-------+----+----------+----------+--------------+---+
|ID     |bool|date      |timestamp |bool_in_2weeks|D  |
+-------+----+----------+----------+--------------+---+
|1092829|0   |2019-10-29|1572307200|0             |0  |
|1092829|1   |2019-12-14|1576281600|1             |0  |
|1092829|1   |2020-07-15|1594771200|2             |1  |
|1092829|1   |2020-07-17|1594944000|1             |0  |
|1092829|0   |2020-07-19|1595116800|0             |0  |
|1092829|1   |2020-08-15|1597449600|1             |0  |
|1092829|1   |2020-09-10|1599696000|1             |0  |
+-------+----+----------+----------+--------------+---+

I took your example and ran a few tests. This is what I came up with:

df = spark.createDataFrame(
    [
        {"ID": 1092829, "date": "2019-10-29", "bool": 0},
        {"ID": 1092829, "date": "2019-12-14", "bool": 1},
        {"ID": 1092829, "date": "2020-07-15", "bool": 1},
        {"ID": 1092829, "date": "2020-07-17", "bool": 1},
        {"ID": 1092829, "date": "2020-07-19", "bool": 0},
        {"ID": 1092829, "date": "2020-08-15", "bool": 1},
        {"ID": 1092829, "date": "2020-09-10", "bool": 1},
    ]
)

For your case, it is both simpler and more adequate to use Window functions.
The window is constructed by partitioning your DataFrame according to your "ID" column. This will guarantee that no two different "ID" values will get their date/bool mixed up when applying the lag. Within the window, you need to sort the date in ascending order for the datediff operation.

from pyspark.sql.window import Window
import pyspark.sql.functions as F

window = Window.partitionBy("ID").orderBy(F.col("ID").asc(), F.col("date").asc())

transformed_df = df.withColumn(
    "D",
    F.coalesce(
        # Apply lag on "date" by -1 (bring next row's value to current row's level) to do `datediff` for the 2-week condition.
        (F.datediff(F.lag("date", -1, None).over(window), F.col("date")) < 14)
        # Apply lag on "bool" by -1 to compare, for a given row, if the next row has a 1 or not.
        & (
            F.lag("bool", -1, None).over(window) == F.lit(1)
        ),  # The comparisons and the "&" produce a boolean column
        F.lit(
            False
        ),  # For the case when the last `date` value for a given "ID" is NULL because of the lag, we set it to False by default.
    ).cast(
        "integer"
    ),  # Cast the boolean result into an integer for a binary value.
)

transformed_df.show()

Resulting DataFrame:

#+-------+----------+----+---+
#|     ID|      date|bool|  D|
#+-------+----------+----+---+
#|1092829|2019-10-29|   0|  0|
#|1092829|2019-12-14|   1|  0|
#|1092829|2020-07-15|   1|  1|
#|1092829|2020-07-17|   1|  0|
#|1092829|2020-07-19|   0|  0|
#|1092829|2020-08-15|   1|  0|
#|1092829|2020-09-10|   1|  0|
#+-------+----------+----+---+

I hope this answers your question.