Spark 3.4中PySpark DataFrame按ID聚合布尔值的最优实现方法

需求回顾 #

现有DataFrame结构如下：

• ID: String类型
• output: Boolean类型

数据示例：

ID      output
AA      true
AA      false
BB      true
BB      true
CC      true
CC      false
CC      true

需要按ID分组，若该ID对应的所有output值均为true则返回true，否则返回false，期望输出：

ID  result
AA  false
BB  true
CC  false

最优实现方案（推荐内置聚合函数） #

直接使用Spark原生聚合函数是性能最优的方案，无需自定义UDF或窗口函数，Spark会对这类操作做全量优化：

方法1：利用bool_and聚合函数 #

bool_and是Spark SQL专门用于逻辑与的聚合函数，会对分组内的所有布尔值做逻辑与运算，只有全部为true时结果才为true：

from pyspark.sql import functions as F

result_df = df.groupBy("ID").agg(F.expr("bool_and(output)").alias("result"))

方法2：利用min函数 #

在Spark中布尔值的排序规则为true > false，因此分组内取min(output)，若结果为true则说明所有值都是true，否则存在false：

from pyspark.sql import functions as F

result_df = df.groupBy("ID").agg(F.min("output").alias("result"))

窗口函数实现方式 #

如果业务场景需要保留原DataFrame的所有行（而非仅分组结果），可以用窗口函数，之后去重得到分组结果：

from pyspark.sql import Window
from pyspark.sql import functions as F

window_spec = Window.partitionBy("ID")
result_df = df.withColumn("result", F.min("output").over(window_spec)) \
              .select("ID", "result") \
              .distinct()

注意：窗口函数会为每一行计算结果，最后需要去重操作，性能略逊于直接groupBy聚合。

UDF实现方式（不推荐） #

UDF需要将数据序列化到Python端处理，性能远不如内置函数，仅在特殊场景下使用：

from pyspark.sql import functions as F
from pyspark.sql.types import BooleanType

def all_true(arr):
    return all(arr)

all_true_udf = F.udf(all_true, BooleanType())
result_df = df.groupBy("ID").agg(all_true_udf(F.collect_list("output")).alias("result"))

方案对比 #

• 最优选择：bool_and或min的groupBy聚合方案，代码简洁且性能最优，Spark原生优化支持。
• 窗口函数：适合需要保留原行数据的场景，但需额外去重操作，性能稍差。
• UDF：性能最差，尽量避免使用。

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