Python pyspark中的熊猫UDF

我试图在spark数据框上填充一系列观察结果。基本上，我有一个天数列表，我应该为每组创建缺少的天数。
熊猫中有

reindex

功能，pyspark中没有该功能。
我尝试实现一个UDF：

@pandas_udf(schema, functionType=PandasUDFType.GROUPED_MAP)
def reindex_by_date(df):
    df = df.set_index('dates')
    dates = pd.date_range(df.index.min(), df.index.max())
    return df.reindex(dates, fill_value=0).ffill()

这看起来应该做我需要的事情，但是它失败了

AttributeError:只能使用具有datetimelike值的.dt访问器

. 我做错了什么？
以下是完整的代码：

data = spark.createDataFrame(
        [(1, "2020-01-01", 0), 
        (1, "2020-01-03", 42), 
        (2, "2020-01-01", -1), 
        (2, "2020-01-03", -2)],
        ('id', 'dates', 'value'))

data = data.withColumn('dates', col('dates').cast("date"))

schema = StructType([
     StructField('id', IntegerType()),
     StructField('dates', DateType()),
     StructField('value', DoubleType())])

@pandas_udf(schema, functionType=PandasUDFType.GROUPED_MAP)
def reindex_by_date(df):
     df = df.set_index('dates')
     dates = pd.date_range(df.index.min(), df.index.max())
     return df.reindex(dates, fill_value=0).ffill()

data = data.groupby('id').apply(reindex_by_date)

理想情况下，我想要这样的东西：

+---+----------+-----+                                                          
| id|     dates|value|
+---+----------+-----+
|  1|2020-01-01|    0|
|  1|2020-01-02|    0|
|  1|2020-01-03|   42|
|  2|2020-01-01|   -1|
|  2|2020-01-02|    0|
|  2|2020-01-03|   -2|
+---+----------+-----+

案例1：每个ID都有一个单独的日期范围。 我会尽量减少udf的内容。在这种情况下，我只计算udf中每个ID的日期范围。对于其他部分，我将使用Spark本机函数

从pyspark.sql导入类型为T
从pyspark.sql导入函数为F
#获取每个ID的最小和最大日期
日期范围=data.groupby（'id'）。agg（F.min（'dates'）。别名（'date\u min'），F.max（'dates'）。别名（'date\u max'））
#计算每个ID的日期范围
@F.udf（returnType=T.ArrayType（T.DateType（）））
def获取日期范围（日期最小值、日期最大值）：
返回列表中t的[t.date（）（pd.date\u范围（date\u min，date\u max））]
#为了每个潜在日期获得一行，我们需要分解UDF输出
日期范围=日期范围。带列(
“日期”，
F.分解（获取日期范围（F.col（'date\u min'）、F.col（'date\u max'））
)
日期范围=日期范围。删除（'date\u min'、'date\u max'）
#为现有条目添加值，为其他条目添加0
结果=日期\u范围。连接(
数据，
['id'，'dates']，
“左”
)
result=result.fillna（{'value'：0}）

案例2：所有ID的日期范围都相同 我认为这里没有必要使用UDF。您想要的内容可以以不同的方式存档：首先，您可以获得所有可能的ID和所有必要的日期。其次，交叉连接它们，这将为您提供所有可能的组合。第三，左键将原始数据连接到组合中。第四，将出现的空值替换为0

#获取所有唯一ID
ids_df=data.select（'id'）.distinct（）
#获取日期序列
date\u min，date\u max=data.agg（F.min（'dates'），F.max（'dates'））.collect（）[0]
dates=[[t.date（）]表示列表中的t（pd.date_范围（date_min，date_max））]
dates\u df=spark.createDataFrame（data=dates，schema=“dates:date”）
#计算所有组合
所有通信=ID交叉连接（日期）
#添加值列
结果=所有_comdinations.join(
数据，
['id'，'dates']，
“左”
)
#将所有空值替换为0
result=result.fillna（{'value'：0}）

请注意此解决方案的以下限制：

交叉连接可能非常昂贵。解决这一问题的一个可能办法可以在中找到

collect语句和Pandas的使用导致了一个不完全并行的Spark转换

---

---

[编辑]分为两种情况，因为我最初认为所有id都有相同的日期范围。

这可能有效，但是

min

和

max

日期对于每个

id

可能不同。在本例中，它们是相同的，但并不总是如此。感谢您的澄清。对于每个

id

都有自己的日期范围的情况，我在答案中添加了另一个解决方案。