Как рассчитать разницу в днях с указанными условиями между строками в pyspark

from pyspark.sql import Window, functions as F

df = spark_session.createDataFrame([
    ('A2A', 'I', date(2023, 4, 1), date(2023, 4, 5)),
    ('A2A', 'D', date(2023, 4, 5), date(2023, 4, 7)),
    ('A2A', 'D', date(2023, 5, 10), date(2023, 5, 13)),
    ('BB3', 'I', date(2023, 5, 5), date(2023, 5, 6)),
    ('BB3', 'I', date(2023, 7, 29), date(2023, 8, 2)),
    ('BB3', 'D', date(2023, 9, 30), date(2023, 10, 3))
],
('ID', 'Type', 'first_date', 'last_date'))
df.show()

# +---+----+----------+----------+
# | ID|Type|first_date| last_date|
# +---+----+----------+----------+
# |A2A|   I|2023-04-01|2023-04-05|
# |A2A|   D|2023-04-05|2023-04-07|
# |A2A|   D|2023-05-10|2023-05-13|
# |BB3|   I|2023-05-05|2023-05-06|
# |BB3|   I|2023-07-29|2023-08-02|
# |BB3|   D|2023-09-30|2023-10-03|
# +---+----+----------+----------+

Теперь нам нужно разбить df на 2 df, чтобы разделить типы I и D.

df_l = df.filter(F.col('Type') == 'I')
df_r = df.filter(F.col('Type') == 'D')

# the next 2 rows are to avoid a known issue with joining dataframes derived from the same one.

df_l = df_l.toDF(*df_l.columns)
df_r = df_r.toDF(*df_r.columns)

Теперь мы соединяем эти два фрейма данных, чтобы иметь возможность выполнить вычитание дат.

df_int = df_l.join(df_r, (df_l['ID'] == df_r['ID']) & (df_l['last_date'] <= df_r['first_date']), 'inner')
df_int.show()

# +---+----+----------+----------+---+----+----------+----------+
# | ID|Type|first_date| last_date| ID|Type|first_date| last_date|
# +---+----+----------+----------+---+----+----------+----------+
# |BB3|   I|2023-07-29|2023-08-02|BB3|   D|2023-09-30|2023-10-03|
# |BB3|   I|2023-05-05|2023-05-06|BB3|   D|2023-09-30|2023-10-03|
# |A2A|   I|2023-04-01|2023-04-05|A2A|   D|2023-04-05|2023-04-07|
# |A2A|   I|2023-04-01|2023-04-05|A2A|   D|2023-05-10|2023-05-13|
# +---+----+----------+----------+---+----+----------+----------+

# Now filter out rows with preceding unnecessary 'I' entries
w = Window.partitionBy(df_l['ID'], df_r['first_date'], df_r['last_date']).orderBy(df_l['first_date'].desc(), df_l['last_date'].desc())
df_int = df_int.withColumn('rn', F.row_number().over(w))
df_int = df_int.filter(F.col('rn') == 1).drop(F.col('rn'))
df_int.show()

# +---+----+----------+----------+---+----+----------+----------+
# | ID|Type|first_date| last_date| ID|Type|first_date| last_date|
# +---+----+----------+----------+---+----+----------+----------+
# |BB3|   I|2023-07-29|2023-08-02|BB3|   D|2023-09-30|2023-10-03|
# |A2A|   I|2023-04-01|2023-04-05|A2A|   D|2023-04-05|2023-04-07|
# |A2A|   I|2023-04-01|2023-04-05|A2A|   D|2023-05-10|2023-05-13|
# +---+----+----------+----------+---+----+----------+----------+

# Adding date difference
f_int = df_int.withColumn('days_between', F.datediff(df_r['first_date'], df_l['last_date']))
df_int.show()

# +---+----+----------+----------+---+----+----------+----------+------------+
# | ID|Type|first_date| last_date| ID|Type|first_date| last_date|days_between|
# +---+----+----------+----------+---+----+----------+----------+------------+
# |BB3|   I|2023-07-29|2023-08-02|BB3|   D|2023-09-30|2023-10-03|          59|
# |A2A|   I|2023-04-01|2023-04-05|A2A|   D|2023-04-05|2023-04-07|           0|
# |A2A|   I|2023-04-01|2023-04-05|A2A|   D|2023-05-10|2023-05-13|          35|
# +---+----+----------+----------+---+----+----------+----------+------------+

# getting rid of unnecessary columns
df_int = df_int.select(df_l['ID'], df_r['Type'], df_r['first_date'], df_r['last_date'], F.col('days_between'))

Теперь нам нужно объединить эту информацию обратно в наш исходный фрейм данных. Мы сделаем левое соединение для сложного ключа, состоящего из ID, Type, first_date, last_date, предполагая, что пары (first_date, last_date) уникальны.

df_res = df.join(df_int, ['ID', 'Type', 'first_date', 'last_date'], 'left').orderBy('ID', 'first_date')
df_res.show()

# +---+----+----------+----------+------------+
# | ID|Type|first_date| last_date|days_between|
# +---+----+----------+----------+------------+
# |A2A|   I|2023-04-01|2023-04-05|        null|
# |A2A|   D|2023-04-05|2023-04-07|           0|
# |A2A|   D|2023-05-10|2023-05-13|          35|
# |BB3|   I|2023-05-05|2023-05-06|        null|
# |BB3|   I|2023-07-29|2023-08-02|        null|
# |BB3|   D|2023-09-30|2023-10-03|          59|
# +---+----+----------+----------+------------+