【问题标题】：pyspark Fill missing data per id in spark dataframepyspark 填充 spark 数据框中每个 id 的缺失数据
【发布时间】：2021-03-11 20:25:57
【问题描述】：

目前我正在使用 pandas 进行一些转换，但我想在 Pyspark 中进行。

我正在使用 Pyspark 3.0.1

这些转换之一是填充每个 id 缺少的日期

我创建了一个假数据集来复制转换：

df = pd.DataFrame(
{
    'id': [1,1,1,2,2,2,3,3,3]
    ,'date':['2020-11-01','2020-11-03','2020-11-05','2020-11-02','2020-11-03','2020-11-05','2020-11-06','2020-11-08'
             ,'2020-11-10']
    ,'amount':[120,200,300,50,90,80,87,67,400]
    ,'fee': [0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9]
}
)

df_disb_date = pd.DataFrame({
    'id': [1,2,3]
    ,'disb_date': ['2020-11-01','2020-11-02','2020-11-07']
})

df

    id  date    amount  fee
0   1   2020-11-01  120 0.1
1   1   2020-11-03  200 0.2
2   1   2020-11-05  300 0.3
3   2   2020-11-02  50  0.4
4   2   2020-11-03  90  0.5
5   2   2020-11-05  80  0.6
6   3   2020-11-06  87  0.7
7   3   2020-11-08  67  0.8
8   3   2020-11-10  400 0.9


df_disb_date

    id  disb_date
0   1   2020-11-01
1   2   2020-11-02
2   3   2020-11-07

转换代码为：

df['date'] = pd.to_datetime(df['date'])

df_max_min_dates = df.groupby(['id'])
df_max_min_dates_grouped =  pd.DataFrame({'max_date': df_max_min_dates['date'].max(),
                                            'min_date': df_max_min_dates['date'].min()})
df_max_min_dates_grouped = df_max_min_dates_grouped.reset_index()
date_max = df_max_min_dates_grouped['max_date'].max().strftime('%Y-%m-%d')
date_min = df_max_min_dates_grouped['min_date'].min().strftime('%Y-%m-%d')
idx = pd.date_range(date_min, date_max)
_id = df['id'].unique()
df = df.set_index(['id','date'])
idx = pd.MultiIndex.from_product((idx, _id), names=['date', 'id'])
df = df.reset_index()
df = df.set_index(['date', 'id']).reindex(idx, fill_value=0).reset_index()
df = df.reset_index()
df = df.merge(df_max_min_dates_grouped[['id','max_date','min_date']], left_on='id',right_on='id').drop(columns = {'max_date','min_date'})


df_disb_date['disb_date'] = pd.to_datetime(df_disb_date['disb_date'])
df = df.merge(df_disb_date,left_on = 'id',right_on = 'id', how = 'left')
df = df[df['date']>=df['disb_date']]
df


    index   date    id  amount  fee disb_date
0   0   2020-11-01  1   120 0.1 2020-11-01
1   3   2020-11-02  1   0   0.0 2020-11-01
2   6   2020-11-03  1   200 0.2 2020-11-01
3   9   2020-11-04  1   0   0.0 2020-11-01
4   12  2020-11-05  1   300 0.3 2020-11-01
5   15  2020-11-06  1   0   0.0 2020-11-01
6   18  2020-11-07  1   0   0.0 2020-11-01
7   21  2020-11-08  1   0   0.0 2020-11-01
8   24  2020-11-09  1   0   0.0 2020-11-01
9   27  2020-11-10  1   0   0.0 2020-11-01
11  4   2020-11-02  2   50  0.4 2020-11-02
12  7   2020-11-03  2   90  0.5 2020-11-02
13  10  2020-11-04  2   0   0.0 2020-11-02
14  13  2020-11-05  2   80  0.6 2020-11-02
15  16  2020-11-06  2   0   0.0 2020-11-02
16  19  2020-11-07  2   0   0.0 2020-11-02
17  22  2020-11-08  2   0   0.0 2020-11-02
18  25  2020-11-09  2   0   0.0 2020-11-02
19  28  2020-11-10  2   0   0.0 2020-11-02
26  20  2020-11-07  3   0   0.0 2020-11-07
27  23  2020-11-08  3   67  0.8 2020-11-07
28  26  2020-11-09  3   0   0.0 2020-11-07
29  29  2020-11-10  3   400 0.9 2020-11-07

我不知道如何在pyspark中执行这个操作，我已经创建了数据框：

df_spark = spark.createDataFrame(
    [
          (1,'2020-11-01',120,0.1)
        , (1,'2020-11-03',200,0.2)
        , (1,'2020-11-05',300,0.3)
        , (2,'2020-11-02',50,0.4)
        , (2,'2020-11-03',90,0.5)
        , (2,'2020-11-05',80,0.6)
        , (3,'2020-11-06',87,0.7)
        , (3,'2020-11-08',67,0.8)
        , (3,'2020-11-10',400,0.9)
        
    ],
    ['id', 'date','amount','fee'] 

df_spark.show()

+---+----------+------+---+
| id|      date|amount|fee|
+---+----------+------+---+
|  1|2020-11-01|   120|0.1|
|  1|2020-11-03|   200|0.2|
|  1|2020-11-05|   300|0.3|
|  2|2020-11-02|    50|0.4|
|  2|2020-11-03|    90|0.5|
|  2|2020-11-05|    80|0.6|
|  3|2020-11-06|    87|0.7|
|  3|2020-11-08|    67|0.8|
|  3|2020-11-10|   400|0.9|
+---+----------+------+---+

df_spark_disb_date = spark.createDataFrame(
    [
          (1,'2020-11-01')
        , (2,'2020-11-02')
        , (3,'2020-11-07')
        
    ],
    ['id', 'disb_date'] 
)

df_spark_disb_date.show()

+---+----------+
| id| disb_date|
+---+----------+
|  1|2020-11-01|
|  2|2020-11-02|
|  3|2020-11-07|
+---+----------+

我想先创建一个包含所有日期的数据框：

from pyspark.sql.types import DateType
import pyspark.sql.functions as F
df_spark = df_spark.withColumn("date", col('date').cast(DateType()))

df_dates = df_spark.agg(F.min(df_spark.date).alias('min_date'), F.max(df_spark.date).alias('max_date'))
df_data_range = df_dates.withColumn('dates_range', F.explode(F.expr('sequence(min_date, max_date, interval 1 day)'))).select('dates_range')

df_data_range.show()

+-----------+
|dates_range|
+-----------+
| 2020-11-01|
| 2020-11-02|
| 2020-11-03|
| 2020-11-04|
| 2020-11-05|
| 2020-11-06|
| 2020-11-07|
| 2020-11-08|
| 2020-11-09|
| 2020-11-10|
+-----------+

但是我不知道如何继续，有办法在 pyspark 或 Pandas UDF 中做到这一点？

更新

我曾尝试使用 Pandas UDF 来做这件事，但我只得到每个 id 的开始和结束日期的缺失值，而不是所有日期：

from pyspark.sql.functions import pandas_udf, PandasUDFType
from pyspark.sql.types import StructField,StringType,IntegerType,DoubleType, LongType,StructType

schema = StructType([StructField("id", LongType(), True),
                     StructField("date", StringType(), True),
                StructField("amount", LongType(), True),
                    StructField("fee", DoubleType(), True)])

@pandas_udf(schema, PandasUDFType.GROUPED_MAP)
def pandas_plus_two(df):
    import pandas as pd
    df['date'] = pd.to_datetime(df['date'])
    date_max = df['date'].max().strftime('%Y-%m-%d')
    date_min = df['date'].min().strftime('%Y-%m-%d')
    idx = pd.date_range(date_min, date_max)
    _id = df['id'].unique()
    df = df.set_index(['id','date'])
    idx = pd.MultiIndex.from_product((idx, _id), names=['date', 'id'])
    df = df.reset_index()
    df_result = df.set_index(['date', 'id']).reindex(idx, fill_value=0).reset_index()
    df_result = df_result.reset_index(drop=True)
    df_result['date'] = df_result['date'].dt.strftime('%Y-%m-%d')
    return df_result  # outputs a pandas DataFrame


df_spark.groupBy('id').apply(pandas_plus_two).show()

+---+----------+------+---+
| id|      date|amount|fee|
+---+----------+------+---+
|  1|2020-11-01|   120|0.1|
|  1|2020-11-02|     0|0.0|
|  1|2020-11-03|   200|0.2|
|  1|2020-11-04|     0|0.0|
|  1|2020-11-05|   300|0.3|
|  3|2020-11-06|    87|0.7|
|  3|2020-11-07|     0|0.0|
|  3|2020-11-08|    67|0.8|
|  3|2020-11-09|     0|0.0|
|  3|2020-11-10|   400|0.9|
|  2|2020-11-02|    50|0.4|
|  2|2020-11-03|    90|0.5|
|  2|2020-11-04|     0|0.0|
|  2|2020-11-05|    80|0.6|
+---+----------+------+---+

【问题讨论】：

标签： python dataframe pyspark

【解决方案1】：

IIUC，您可以从 df_spark(dt_max) 中找到max 日期，然后使用sequence 函数获取 disb_date 和 dt_max，将此日期数组分解为行，然后与 df_spark 进行左连接：

from pyspark.sql import functions as F

dt_max = df_spark.select('date').agg(F.max('date')).first()[0]

df_new = df_disb_date.selectExpr(
    "id", 
    "disb_date", 
    f"explode(sequence(date(disb_date), date('{dt_max}'))) as date"
).join(df_spark, ["id","date"], "left") \
.fillna(0, subset=["amount", "fee"])

输出：

df_new.orderBy('id','date').show(100)
+---+----------+----------+------+---+                                          
| id|      date| disb_date|amount|fee|
+---+----------+----------+------+---+
|  1|2020-11-01|2020-11-01|   120|0.1|
|  1|2020-11-02|2020-11-01|     0|0.0|
|  1|2020-11-03|2020-11-01|   200|0.2|
|  1|2020-11-04|2020-11-01|     0|0.0|
|  1|2020-11-05|2020-11-01|   300|0.3|
|  1|2020-11-06|2020-11-01|     0|0.0|
|  1|2020-11-07|2020-11-01|     0|0.0|
|  1|2020-11-08|2020-11-01|     0|0.0|
|  1|2020-11-09|2020-11-01|     0|0.0|
|  1|2020-11-10|2020-11-01|     0|0.0|
|  2|2020-11-02|2020-11-02|    50|0.4|
|  2|2020-11-03|2020-11-02|    90|0.5|
|  2|2020-11-04|2020-11-02|     0|0.0|
|  2|2020-11-05|2020-11-02|    80|0.6|
|  2|2020-11-06|2020-11-02|     0|0.0|
|  2|2020-11-07|2020-11-02|     0|0.0|
|  2|2020-11-08|2020-11-02|     0|0.0|
|  2|2020-11-09|2020-11-02|     0|0.0|
|  2|2020-11-10|2020-11-02|     0|0.0|
|  3|2020-11-07|2020-11-07|     0|0.0|
|  3|2020-11-08|2020-11-07|    67|0.8|
|  3|2020-11-09|2020-11-07|     0|0.0|
|  3|2020-11-10|2020-11-07|   400|0.9|
+---+----------+----------+------+---+

【讨论】：