PySpark百万S前缀列高效处理：生成非空列名列表方案咨询

我有如下PySpark DataFrame：

+--------+-------------+---------+---------+---------+
|    code|    updatedAt|S0x223433|S1yd33333|S4r256467|
+--------+-------------+---------+---------+---------+
|AAAAAAAA|1713292448319|     3243|     null|     3422|
|BBBBBBBB|1689430451041|     3455|     2345|     7654|
+--------+-------------+---------+---------+---------+

需要为该DataFrame新增一列，每行仅保留非空的S前缀列的列名，结果DataFrame如下：

+--------+-------------+---------+---------+---------+---------------------+
|    code|    updatedAt|S0x223433|S1yd33333|S4r256467|new_col             |
+--------+-------------+---------+---------+---------+---------------------+
|AAAAAAAA|1713292448319|     3243|     null|     3422|[S0x223433,S4r256467]|
|BBBBBBBB|1689430451041|     null|     2345|     7654|[S1yd33333,S4r256467]|
+--------+-------------+---------+---------+---------+---------------------+

需求细节 #

DataFrame拥有近百万列，期望处理逻辑：

1. 检查列名是否匹配正则表达式（匹配以S开头的列）
2. 若该列值非空，则将列名加入列表
3. 返回每行对应的符合条件的列名列表
   （后续会将新增列中的列表展开以实现数据规范化）

当前问题 #

当前新增列的逻辑因遍历百万列运行极慢：

non_null_s_columns= array([when(col(c).isNotNull(), lit(c)) for c in SPrefixedcolumns])

已尝试方案 #

• 删除空列：因数据集稀疏无济于事
• 编写UDF接收行数据生成新列：尚未成功

求助：是否存在更高效的实现方式？希望获取PySpark DataFrame或Glue动态帧相关的实现思路，避免遍历百万列。

方案1：使用Spark内置函数map_from_entries结合stack（推荐） #

当列数极多时，stack可以将宽表转为长表，再进行聚合，避免遍历所有列生成表达式：

1. 筛选S前缀列：先获取所有匹配正则的列名

import re
from pyspark.sql import functions as F

# 匹配以S开头的列
s_cols = [c for c in df.columns if re.match(r'^S', c)]

2. 构造stack表达式：将所有S列转为(key, value)行

# 构造stack的参数：列数 + 每个(key, value)对
stack_expr = f"stack({len(s_cols)}, {', '.join([f'{repr(c)}, {c}' for c in s_cols])}) as (col_name, col_value)"

3. 长表聚合生成列表：过滤空值后，按原主键分组聚合列名

# 假设主键是code和updatedAt，根据实际情况调整
result = df.select("code", "updatedAt", F.expr(stack_expr)) \
           .filter(F.col("col_value").isNotNull()) \
           .groupBy("code", "updatedAt") \
           .agg(F.collect_list("col_name").alias("new_col")) \
           .join(df, on=["code", "updatedAt"], how="right")

这种方式的优势是利用Spark的分布式计算处理长表，避免生成百万个when表达式，适合超宽表场景。

方案2：利用Glue动态帧的relationalize方法（针对AWS Glue场景） #

如果使用AWS Glue，可以直接用relationalize将宽表转为规范化的长表，天然满足后续展开需求：

from awsglue.dynamicframe import DynamicFrame

# 将DataFrame转为DynamicFrame
dynamic_df = DynamicFrame.fromDF(df, glue_context, "wide_df")

# 生成规范化的表，参数index设定主键列
normalized_dyf = dynamic_df.relationalize("root", "s3://your-bucket/temp_path", index=["code", "updatedAt"])

# 获取包含列名和值的子表，过滤空值后即可得到所需结果
s_prefix_dyf = normalized_dyf.select_fields(["code", "updatedAt", "col_name", "col_value"]) \
                             .filter(F.col("col_value").isNotNull())

relationalize会自动处理超宽表的转换，适合Glue环境下的稀疏宽表场景，后续直接可以基于长表做展开，无需额外生成列表列。

方案3：优化原表达式生成逻辑（若必须保留宽表结构） #

如果一定要保留原宽表结构，优化表达式生成的效率：

• 避免在列表推导中重复调用col()和lit()，可以预生成表达式
• 利用Spark的array_remove过滤null值

# 生成表达式列表时，用列表推导提前构建
exprs = [F.when(F.col(c).isNotNull(), F.lit(c)) for c in s_cols]
# 用array_remove去掉null元素
non_null_s_columns = F.array_remove(F.array(*exprs), None)

result = df.withColumn("new_col", non_null_s_columns)

相比原代码，array_remove可以减少后续的空值处理开销，但本质还是生成大量表达式，对于百万列场景，性能提升有限，不如方案1高效。

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