Pyspark如何根据DataFrame的ID列取值获取对应threshold阈值

你之前的写法threshold = df.first()['ID']只能拿到DataFrame第一行的ID值，无法获取全量ID对应的阈值映射，所以需要先对ID做去重再收集所有映射关系。

方案1：生成本地查询字典（适合在Driver端按ID取阈值的场景） #

你的数据只有数千条，完全可以将去重后的ID与阈值映射关系收集到Driver端转成Python字典，查询效率极高：

# 按ID去重后提取ID、percentage两列，收集后转成字典
id_threshold_map = dict(df.dropDuplicates(["ID"]).select("ID", "percentage").collect())

使用示例：

# 获取B101对应的阈值
threshold = id_threshold_map["B101"]
# 获取B202对应的阈值
threshold = id_threshold_map["B202"]

方案2：DataFrame新增threshold列（适合后续分布式计算场景） #

如果你需要在Spark的分布式计算逻辑中用到每个ID对应的阈值，不需要拉取数据到本地，直接用窗口函数新增列即可：

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

# 按ID分区，同一个ID的threshold取对应percentage的值（同ID的percentage一致，用max/min/first都可以）
id_window = Window.partitionBy("ID")
df = df.withColumn("threshold", F.first("percentage").over(id_window))

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