PySpark查找仅含1与Null的DataFrame列的技术求助

你的代码无法得到正确结果，原因是df(column).isin(1)生成的是一个布尔类型的列对象，而非判断该列所有值是否仅包含1和Null，Python会将这个列对象视为真值，最终会选中所有列。

下面是正确的实现方式：

方法1：通过收集每列唯一值筛选（适合小数据集） #

# 遍历所有列，收集每列的唯一值集合
col_unique_vals = {col: set(df.select(col).distinct().rdd.flatMap(lambda x: x).collect()) for col in df.columns}

# 筛选出仅包含1和Null的列（排除全Null的列，确保列中存在1）
target_columns = [col for col, vals in col_unique_vals.items() if vals.issubset({1, None}) and 1 in vals]

# 选择目标列并展示
df.select(target_columns).show()

方法2：通过Spark内置函数聚合判断（适合大数据集） #

如果数据集规模较大，把所有列的唯一值拉到Driver节点可能导致内存不足，推荐用分布式聚合的方式：

from pyspark.sql import functions as F

# 对每列定义两个聚合判断逻辑：有效不同值数量、是否包含1
agg_exprs = [
    F.countDistinct(F.when(F.col(col).isNull() | (F.col(col) == 1), F.col(col))).alias(f"{col}_valid_distinct"),
    F.count(F.when(F.col(col) == 1, True)).alias(f"{col}_has_1")
    for col in df.columns
]

# 执行聚合计算
agg_result = df.agg(*agg_exprs).collect()[0]

# 筛选符合条件的列：有效不同值只有1（说明仅含1和Null）且存在1
target_columns = [
    col for col in df.columns
    if agg_result[f"{col}_valid_distinct"] == 1 and agg_result[f"{col}_has_1"] > 0
]

# 选择目标列展示
df.select(target_columns).show()

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