Spark能否按列值总和筛选列？PySpark稀疏表列过滤求助

我来帮你搞定这个需求！你已经走对了第一步——用groupBy().sum()拿到列总和，接下来只需要把符合阈值条件的列名筛选出来，再用这些列名去选择原DataFrame的列就行，完全不需要手动列名，几千列也能轻松处理。

完整实现步骤 #

这里以你的示例数据和阈值3为例，给出可扩展的代码：

from pyspark.sql import SparkSession

# 初始化SparkSession（如果还没初始化的话）
spark = SparkSession.builder.appName("FilterColumnsBySum").getOrCreate()

# 模拟你的原始数据
data = [
    (1, 0, 1, 1, 0, 0),
    (1, 1, 0, 0, 0, 0),
    (1, 0, 0, 1, 1, 1),
    (1, 0, 0, 1, 1, 1),
    (1, 1, 0, 0, 1, 0),
    (0, 0, 1, 0, 1, 0)
]
df = spark.createDataFrame(data, ["a", "b", "c", "d", "e", "f"])

# 1. 计算所有列的总和，得到单行DataFrame
sum_df = df.groupBy().sum()

# 2. 将总和结果转换成字典，方便遍历处理
sum_dict = sum_df.collect()[0].asDict()

# 3. 定义阈值，筛选出总和≥阈值的列名
threshold = 3
# 用正则提取原列名（避免列名含特殊字符时截取出错）
import re
selected_cols = []
for sum_col, value in sum_dict.items():
    col_name = re.search(r'sum\((.*?)\)', sum_col).group(1)
    if value >= threshold:
        selected_cols.append(col_name)

# 4. 用筛选后的列名选择原DataFrame的列
filtered_df = df.select(selected_cols)

# 查看结果
filtered_df.show()

代码解释 #

• 步骤1：groupBy().sum()生成一个仅一行的DataFrame，列名格式为sum(原列名)，值对应原列的总和。
• 步骤2：collect()[0]取出这唯一的一行数据，asDict()将其转为字典，键是sum(原列名)，值是列总和。
• 步骤3：用正则表达式提取原列名（适配各种列名格式），再筛选出总和≥阈值的列名列表。
• 步骤4：直接用筛选后的列名列表调用select()，就能得到只保留符合条件列的DataFrame。

可扩展性说明 #

不管你的表有几千列还是更多，这个流程完全自动化：不需要手动列名，不需要硬编码，所有列都会被自动遍历和筛选；而且sum_df只有一行数据，collect()的性能开销可以忽略不计。

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