python - 如何根据每个分区重新索引数据框

Question

假设我有 pyspark 创建的以下数据框

id  date         deleted
1   2019-02-07     true
1   2019-02-04     false
2   2019-02-01     true
3   2019-02-08     false
3   2019-02-06     true

我想从最早的日期到现在(比如2019-02-09)每天重新索引这个表，最早的日期基于每个id，例如，对于id 1，最早的日期是2019- 02-04，对于 id 3，最早日期是 2019-02-06。预期结果是:

id  date         deleted
1   2019-02-04     false
1   2019-02-05     null
1   2019-02-06     null
1   2019-02-07     true
1   2019-02-08     null
1   2019-02-09     null

2   2019-02-01     true
2   2019-02-02     null
      ...
2   2019-02-09     null

3   2019-02-06     true
3   2019-02-07     null
3   2019-02-08     false
3   2019-02-09     null

我知道如何根据所有 ID(即 2019-02-01)获取最早的日期，然后只需构建一个数据框，其中包含每个日期从 2019-02-01 到 2019-02-09 的所有日期id(交叉连接)，然后左连接原始数据框。这种方法的问题是，如果有一个日期说 1980-01-01，那么重新索引将为所有 id 填充从 1980-01-01 到现在的所有数据，这是没有意义的，并且会影响性​​能在此数据帧上执行以下 ETL。

对于基于每个分区的最早日期，没有找到好的方法。

Answer 1

假设您的原始 DataFrame 名为 df，并且 date 列实际上是 DateType 类型:

import pyspark.sql.functions as F
from pyspark.sql.types import DateType, ArrayType
import datetime

# create a UDF to create a range of dates from a start
# date until today
def construct_date_range(start_date):
    ndays = (datetime.datetime.today() - start_date).days
    return reversed([base - datetime.timedelta(days=x) for x in range(0, ndays+1)])
date_range_udf = F.udf(construct_date_range, ArrayType(DateType()))

# find the first date for each id, and create a record for
# all dates since the first
id_dates = (
    df
    .groupBy('id')
    .agg(F.min('date').alias('min_date'))
    .withColumn('date_arr', construct_date_range('min_date'))
    .select('id', F.explode('date_arr').alias('date'))
)

result = id_dates.join(df, on=['id','date'], how='left')