从熊猫到pyspark-将带有开始和结束日期的数据框转换为每日数据？

Question

我有一个使用Pandas的ETL脚本，为了使其更具扩展性，我尝试使用Pyspark重新创建它。到目前为止，一切都很顺利，但是在转换到每日数据集时遇到了问题。每个ID都有一条记录，其中包含开始日期和结束日期

id  age state   start_date  end_date
123 18  CA     2/17/2019    5/4/2019
223 24  AZ     1/17/2019    3/4/2019

我想为开始日期和结束日期之间的每一天创建一条记录，因此我可以将其日常活动数据加入其中。目标输出看起来像这样

id  age state   start_date
123 18  CA      2/17/2019
123 18  CA      2/18/2019
123 18  CA      2/19/2019
123 18  CA      2/20/2019
123 18  CA      2/21/2019
            …
123 18  CA      5/2/2019
123 18  CA      5/3/2019
123 18  CA      5/4/2019

当然要对数据集中的所有ID及其各自的开始日期执行此操作。我可以使用以下方法在Pandas中做到这一点

melt = df.melt(id_vars=['id', 'age', 'state'], value_name='date').drop('variable', axis=1)
melt['date'] = pd.to_datetime(melt['date'])

melt = melt.groupby('id').apply(lambda x: x.set_index('date').resample('d').first())\
           .ffill()\
           .reset_index(level=1)\
           .reset_index(drop=True)

但是我对Pyspark并不陌生（并且在Pandas中为此感到苦恼），所以我被困在这里。非常感谢您的帮助-谢谢！

Answer 1

在this post中找到了解决方案。解决方案的关键是爆炸功能，该功能可以满足我的需求。

解决我的具体示例的代码是

def date_range(t1, t2, step=60*60*24):
    return [t1 + step*x for x in range(int((t2-t1)/step)+1)]

date_range_udf = udf(date_range, ArrayType(LongType()))

df = dataF.select("id",
expr("stack(2, 'start_date', start_date, 'end_date', end_date) as (class_date,date)"))

df_base = \
    df.groupBy('id')\
        .agg(min('date').cast('integer').alias('date_min'), max('date').cast('integer')\
    .alias('date_max'))\
        .withColumn("date", explode(date_range_udf("date_min", "date_max")))\
        .drop('date_min', 'date_max')\
        .withColumn("date", from_unixtime("date"))

哪个给出以下输出（我可以使用它来连接任何其他数据）