main函数return语句提前终止致PySpark+pytest测试失败排查

核心问题就是你判断的——return语句位置错误，当同时传入两个文件时，代码在处理完第一个文件的转换逻辑后就直接return了，后续的合并流程完全没机会执行，自然返回的DataFrame列数不符合预期。

具体修复步骤 #

1. 重构分支逻辑 #

把单文件、双文件的处理流程拆成清晰的分支，确保双文件场景下能完整执行两个文件的转换+合并操作。

2. 统一测试模式的return位置 #

将测试模式下的return语句移到所有业务逻辑的最后，避免在单文件分支里提前终止函数执行。

修正后的示例代码 #

def process_pyspark_data(x_path=None, l_path=None, is_test=False):
    # 初始化最终结果DF
    final_df = None

    # 处理X文件的转换逻辑
    x_transformed_df = None
    if x_path:
        x_df = spark.read.parquet(x_path)
        # 这里替换成你的X文件具体转换逻辑
        x_transformed_df = x_df.withColumnRenamed("prem_old_col", "prem_new_col")

    # 处理L文件的转换逻辑
    l_transformed_df = None
    if l_path:
        l_df = spark.read.parquet(l_path)
        # 这里替换成你的L文件具体转换逻辑
        l_transformed_df = l_df.withColumnRenamed("loss_old_col", "loss_new_col")

    # 单文件场景处理
    if x_path and not l_path:
        final_df = x_transformed_df
        if not is_test:
            final_df.write.mode("overwrite").parquet("output/x_single_output")
    elif l_path and not x_path:
        final_df = l_transformed_df
        if not is_test:
            final_df.write.mode("overwrite").parquet("output/l_single_output")
    # 双文件合并场景处理
    elif x_path and l_path:
        # 替换成你的实际合并逻辑，比如unionByName
        final_df = x_transformed_df.unionByName(l_transformed_df, allowMissingColumns=False)
        if not is_test:
            final_df.write.mode("overwrite").parquet("output/combined_output")

    # 测试模式统一返回最终处理后的DF
    if is_test:
        return final_df

关键改动说明 #

• 把测试用的return放在函数末尾，确保所有转换、合并逻辑执行完毕后才返回结果。
• 双文件场景单独分支处理，保证两个文件的转换逻辑都执行完成后再合并。
• 单文件场景只在没有另一个文件时才赋值最终DF，避免覆盖双文件的合并结果。

这样修改后，双文件测试时就能完整走完合并流程，返回的DataFrame列数会符合预期的26列。

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