如何使用PySpark UDF按日期范围将患者记录按年拆分为多行

原有代码问题梳理 #

• 变量名拼写错误：函数返回时调用了未定义的last_date_of_year，实际应为提前定义的last_date_in_year
• 逻辑错误：return语句放在while循环内部，只会返回第一个年份的拆分结果，无法返回所有年份的记录
• UDF类型不匹配：普通UDF仅支持1条输入对应1条输出，你的需求是1条输入对应多条输出，需要返回数组结构后配合explode函数拆分行
• 返回类型定义错误：原代码指定返回StringType()，但实际返回多值结构，需自定义匹配的Struct类型

修正后完整可运行代码 #

import pyspark.sql.functions as f
from pyspark.sql.types import *

# 修正后的处理函数：返回所有拆分后的记录列表
def process_data(identifier, eff_date, end_date):
    result = []
    eff_year = int(eff_date.split("/")[2])
    end_year = int(end_date.split("/")[2])
    
    current_year = eff_year
    while current_year <= end_year:
        if str(current_year) in eff_date:
            first_date_in_year = eff_date
        else:
            first_date_in_year = "%02d/%02d/%4d" % (1, 1, current_year)
            
        if str(current_year) in end_date:
            last_date_in_year = end_date
        else:
            last_date_in_year = "%02d/%02d/%4d" % (12, 31, current_year)
            
        # 追加当前年份的记录到结果列表
        result.append((identifier, first_date_in_year, last_date_in_year))
        current_year += 1
    return result

# 定义UDF返回类型：数组嵌套结构，对应多条拆分记录
udf_return_type = ArrayType(
    StructType([
        StructField("Identifier", StringType(), nullable=False),
        StructField("New_Eff_Date", StringType(), nullable=False),
        StructField("New_End_Date", StringType(), nullable=False)
    ])
)

# 注册UDF
process_data_udf = f.udf(process_data, udf_return_type)

# 调用UDF并拆分行得到最终结果
df_explode = df.withColumn("split_records", process_data_udf("Identifier", "Eff_Date", "End_Date")) \
               .select(f.explode("split_records").alias("record")) \
               .select("record.*")

# 查看结果
df_explode.show(truncate=False)

运行说明 #

上述代码执行后df_explode的数据格式完全匹配你提供的期望输出，无需额外调整。如果你的Spark版本在3.0以上，也可以使用Pandas UDTF实现，性能会比普通Python UDF更好。

内容的提问来源于stack exchange，提问作者Bhavani Kumar Metla