PySpark条件判断代码优化求助：避免重复执行相同逻辑

你的代码中三个分支存在大量重复逻辑，仅性别过滤条件不同，可以通过提取公共逻辑+动态构造过滤条件的方式优化，以下是两种可行方案：

方案一：分步构造过滤条件 #

先定义所有分支都需要的基础过滤条件，再根据sexe参数动态添加性别过滤规则，最后统一执行后续的数据处理操作：

import pyspark.sql.functions as F

# 基础过滤条件：所有分支都需要的column_flag判断
base_filter = F.col("column_flag") == False

# 根据sexe参数追加性别过滤条件
if sexe == "male":
    final_filter = base_filter & (F.col("sexe") == 1)
elif sexe == "female":
    final_filter = base_filter & (F.col("sexe") == 2)
else:
    final_filter = base_filter

# 统一执行数据处理，避免重复代码
new_df = df.where(final_filter).withColumn("new_column", F.col("column1") / F.col("column3"))

方案二：字典映射简化条件判断 #

用字典建立sexe参数到对应过滤条件的映射，后续新增性别类型只需修改字典，无需添加新的分支判断，扩展性更强：

import pyspark.sql.functions as F

# 性别参数与过滤条件的映射表
sexe_filter_map = {
    "male": F.col("sexe") == 1,
    "female": F.col("sexe") == 2
}

# 组合最终过滤条件
final_filter = F.col("column_flag") == False
# 如果sexe在映射表中，追加对应过滤条件
if sexe in sexe_filter_map:
    final_filter = final_filter & sexe_filter_map[sexe]

# 统一执行数据处理
new_df = df.where(final_filter).withColumn("new_column", F.col("column1") / F.col("column3"))

两种方案都能彻底消除代码冗余，后续修改公共逻辑（比如调整new_column的计算规则、修改column_flag的判断条件）时，只需修改一处即可，大幅提升代码的可维护性。

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