如何在PySpark DataFrame中按分组判断指定值是否存在并生成布尔标记

你需要的功能可以通过窗口函数+条件聚合实现，和pandas的groupby().transform()效果完全一致，完整可运行代码如下：

from pyspark.sql import functions as f
from pyspark.sql.window import Window

# 示例df构造代码你已经有了，这里直接写核心逻辑
# 1. 定义按id分组的窗口
id_window = Window.partitionBy("id")

# 2. 计算flag字段
df = df.withColumn(
    "flag",
    # 同组内只要有一个fruit为pear，最大值就为True
    f.max(f.when(f.col("fruit") == "pear", True).otherwise(False)).over(id_window)
)

# 打印结果验证
df.show()

运行后输出和你预期的完全一致：

+---+-------+------+
| id|  fruit|  flag|
+---+-------+------+
|  a|  apple|  true|
|  a|   pear|  true|
|  b|   pear|  true|
|  c| carrot| false|
|  c|  apple| false|
+---+-------+------+

逻辑说明 #

• Window.partitionBy("id")的作用等价于pandas中按id分组，窗口内包含同一个id的所有行
• when(f.col("fruit") == "pear", True).otherwise(False)逐行判断当前行的fruit值是否为pear，返回布尔值
• max(布尔值).over(id_window)取同组内布尔值的最大值：只要有一个值为True，全组的计算结果就为True，刚好匹配你“分组下至少存在一个pear则全组flag为True”的需求
• 最终生成的flag字段原生为布尔类型，不需要额外转换

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