PySpark如何从其他DataFrame提取列名列表替代手动列定义

核心逻辑 #

直接通过DataFrame内置的columns属性获取parameter表的全部列名，过滤掉非目标列后自动生成需要的apps_id列表，完全不需要手动维护200+列的清单。

基础实现代码 #

# 动态提取列名：排除首列_c0，剩余列即为需要的apps_id列表
list_selected_apps_id = [col_name for col_name in parameter.columns if col_name != '_c0']

# 沿用原有计算逻辑，替换硬编码列表即可
fd_orion_apps = fd_orion_apps.groupBy('msisdn', 'apps_id')\
    .pivot('apps_id')\
    .count()\
    .select('msisdn', *list_selected_apps_id)\
    .fillna(0)
fd_orion_apps.show(truncate=False)

性能优化写法 #

pivot操作默认会先扫描全表统计所有可能的透视列值，数据量大时效率很低，你可以直接把提前拿到的list_selected_apps_id作为第二个参数传入pivot方法，跳过全表扫描步骤，执行速度会明显提升：

list_selected_apps_id = [col_name for col_name in parameter.columns if col_name != '_c0']

fd_orion_apps = fd_orion_apps.groupBy('msisdn', 'apps_id')\
    .pivot('apps_id', list_selected_apps_id)  # 提前指定透视列值，跳过全表扫描
    .count()\
    .select('msisdn', *list_selected_apps_id)\
    .fillna(0)
fd_orion_apps.show(truncate=False)

注意事项 #

• 如果你需要保留parameter表的_c0列参与计算，直接删除列表推导式里的if col_name != '_c0'过滤条件，使用parameter.columns即可拿到全部列名
• 动态提取的列名和pivot生成的列名类型完全一致，不会出现硬编码时可能出现的类型不匹配、列名写错的问题

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