Databricks中PySpark筛选非全True布尔列的简便方法

问题背景 #

我在Databricks环境中使用PySpark创建了如下DataFrame：

columns = ["a", "b", "c"]
data = [(True, True, True), (True, True, True), (True, False, True)]
df = spark.createDataFrame(data).toDF(*columns)
df.display()

我的需求是筛选出所有值并非全为True的列。在Pandas中可以用df['a'].all()判断某列是否全为True，但PySpark里没找到直接等效的方法。目前我用了一个复杂的解决方案：

df.select(*[column for column in df.columns 
            if df.select(column).distinct().collect() != 
            spark.createDataFrame([True], 'boolean').toDF(column).collect()])

这个方法能得到预期结果，但想知道PySpark里有没有更简便的实现方式？

简便实现方案 #

方法1：利用聚合函数min()判断 #

布尔类型中False等价于0，True等价于1，若某列全为True，min(column)的结果就是True，反之则为False。我们可以先计算各列的最小值，再筛选出最小值为False的列：

from pyspark.sql.functions import min

# 计算每列的布尔最小值
agg_result = df.agg(*[min(c).alias(c) for c in df.columns]).collect()[0]
# 筛选出非全True的列名
non_all_true_cols = [col_name for col_name in agg_result if not agg_result[col_name]]
# 生成结果DataFrame
result_df = df.select(*non_all_true_cols)
result_df.display()

方法2：通过filter+count判断列中是否存在False #

直接判断每列是否存在False，如果filter(~col(c)).count() > 0，说明该列不是全True：

from pyspark.sql.functions import col

non_all_true_cols = [c for c in df.columns if df.filter(~col(c)).count() > 0]
result_df = df.select(*non_all_true_cols)
result_df.display()

方法3：Spark 3.0+ 使用exists函数 #

Spark 3.0及以上版本支持exists函数，可直接判断列中是否存在不符合条件的行：

from pyspark.sql.functions import exists, col

# 遍历列，判断每列是否存在False
non_all_true_cols = []
for c in df.columns:
    has_false = df.select(exists(df, lambda row: ~row[c])).collect()[0][0]
    if has_false:
        non_all_true_cols.append(c)

result_df = df.select(*non_all_true_cols)
result_df.display()

这些方案都比原实现更简洁高效，其中方法1和方法2在大数据场景下性能更优，避免了不必要的DataFrame创建和多次collect操作。

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