PySpark/Python：基于可变列输入创建计算列的最优方案

数据结构说明 #

• age_group{#}：符合用户定义年龄组的记录数，最多支持10个年龄组（如age_group1为0-10岁记录数）。
• bin_number：用户可定义多组年龄分组规则（如bin1的age_group1为0-10，bin2的age_group1为0-17）。
• t_group：按县统计的对应年龄组总人口，列名格式为t_年龄组编号_bin编号（如t_group1_1为bin1中0-10岁总人口）。

样本数据 #

（注：原样本数据中t_group2_2列重复，已修正为t_group1_2以匹配列名规则）

需求目标 #

新增计算列，规则如下：

• 当bin_number=1时，计算(age_group1/t_group1_1)*100000、(age_group2/t_group2_1)*100000
• 当bin_number=2时，计算(age_group1/t_group1_2)*100000、(age_group2/t_group2_2)*100000

现有方案及问题 #

目前采用循环遍历bin值、过滤后计算再union的方式，但代码硬编码了bin数量，扩展性差。尝试过groupBy、窗口函数未找到可行方案，希望实现更简洁高效的写法，且最终输出保持原表结构。

现有代码 #

from functools import reduce
import pyspark.sql.functions as F

RATE = 100000
i = 1
dfs = []
while i <= 2:
    dfGroup = df.filter(F.col("bin_number") == i)  # df为目标表
    
    totalBins = [x for x in df.columns if x.startswith("t_group") and x.endswith(str(i))]
    dfGroup = dfGroup.select(
        "*",
        *[((F.col(f"age_group{x}") / F.col(f"t_group{x}_{i}")) * RATE).alias(
                f"crude_rate_age_group_bin_{x}"
            ) for x in range(1, len(totalBins)+1)],
    )
   
    dfs.append(dfGroup)

    i += 1  # 原代码i += i会导致死循环，此处修正为i += 1

dfRate = reduce(F.DataFrame.unionAll, dfs)

优化方案 #

通过动态生成条件计算列的方式避免循环和硬编码，直接基于原表一次性生成所有需要的计算列，无需拆分再union，效率更高且扩展性强：

import pyspark.sql.functions as F

RATE = 100000

# 自动提取所有年龄组编号
age_group_cols = [col for col in df.columns if col.startswith("age_group")]
age_group_nums = [col.split("_")[-1] for col in age_group_cols]

# 动态生成每个年龄组的计算列
calc_cols = []
for num in age_group_nums:
    # 根据bin_number的值动态匹配对应的t_group列
    t_group_col = F.when(F.col("bin_number") == 1, F.col(f"t_group{num}_1")) \
                   .when(F.col("bin_number") == 2, F.col(f"t_group{num}_2")) \
                   .otherwise(F.lit(None))
    # 计算粗率并命名
    crude_rate = ((F.col(f"age_group{num}") / t_group_col) * RATE).alias(f"crude_rate_age_group_bin_{num}")
    calc_cols.append(crude_rate)

# 拼接原表列与所有计算列
dfRate = df.select("*", *calc_cols)

方案优势 #

1. 无硬编码限制：后续新增bin值，只需在when条件中添加对应分支即可，无需修改核心逻辑
2. 性能更优：避免拆分表再union的Shuffle操作，直接在原表上完成计算
3. 扩展性强：年龄组数量变化时，代码会自动识别所有age_group列，无需手动调整
4. 代码更简洁：逻辑清晰，减少冗余代码量

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