PySpark Pandas UDF合并多行日期范围缺失id列问题求解

问题根因 #

你当前使用的merge_dates函数返回的Pandas DataFrame仅包含start_dt和end_dt两列，与指定的3列schema（id、start_dt、end_dt）不匹配，因此触发列数不匹配报错；修改schema后虽然能正常运行，但自然缺失了id字段。

解决方法 #

仅需修改merge_dates函数，在返回合并结果前补充分组对应的id字段即可，每个分组内所有行的id完全一致，直接取第一行的id赋值即可。

修改后的聚合函数代码 #

from datetime import datetime

mydt = datetime(1970,1,1).date()
def merge_dates(grp):
    # 取当前分组的id，同一分组内所有行id相同
    current_id = grp["id"].iloc[0]
    dt_groups = ((grp["start_dt"] - grp["end_dt"].shift(fill_value=mydt)).dt.days > 1).cumsum()
    grouped = grp.groupby(dt_groups).agg({"start_dt":"min", "end_dt":"max"}).reset_index(drop=True)
    # 补充id列到第一列位置
    grouped.insert(0, "id", current_id)
    return grouped if len(grp) == len(grouped) else merge_dates(grouped)

Spark调用代码 #

直接使用原schema即可，无需修改：

data_out = data_in.groupby("id").applyInPandas(merge_dates, schema=data_in.schema)
data_out.show()

输出结果 #

运行后将得到和Pandas本地测试完全一致的结果，包含完整的id字段：

+---+----------+----------+
| id| start_dt|   end_dt|
+---+----------+----------+
|  1|2016-12-01|2016-12-02|
|  1|2016-12-05|2017-09-30|
|  1|2017-10-03|2017-10-03|
|  1|2018-05-01|2018-10-31|
|  2|2016-12-01|2016-12-25|
+---+----------+----------+

内容的提问来源于stack exchange，提问作者ironv