如何高效修改PySpark DataFrame重复组内指定行的Cost值？

数据样例 #

给定的PySpark DataFrame deptDF 数据如下：

+--------+----+------------+----------+--------+-----------+
|quantity|cost|participants|activity  |category|id         |
+--------+----+------------+----------+--------+-----------+
|4       |10  |2           |skiing    |outdoor |8589934592 |
|4       |13  |3           |golf      |indoor  |17179869184|
|4       |10  |5           |swimming  |outdoor |25769803776|
|4       |10  |3           |basketball|indoor  |34359738368|
|4       |11  |7           |pool      |indoor  |42949672960|
|4       |11  |12          |pool      |outdoor |51539607552|
|4       |13  |15          |golf      |indoor  |60129542144|
+--------+----+------------+----------+--------+-----------+

需求说明 #

识别出quantity-category-activity组合重复的行，将每组中participants数值更小的行的cost字段设为0。

原实现思路及问题 #

原思路是先通过窗口函数筛选重复行，转换为Python列表后迭代修改原DataFrame，但这种方式需要collect()数据到Driver端，在大数据场景下会导致性能瓶颈，甚至内存溢出。原代码如下：

w = Window.partitionBy('quantity', 'cost', 'activity')
deptDF_duplicates = deptDF.select('*', f.count('quantity').over(w).alias('dupeCount'))\
    .where('dupeCount > 1')\
    .drop('dupeCount')
deptDF_duplicates.show()

duplicates_list = [list(row) for row in deptDF_duplicates["id", "cost"].collect()]

高效的DataFrame原生实现方案 #

可以直接通过PySpark窗口函数和条件函数完成所有操作，全程在分布式环境执行，无需将数据拉取到Driver端，具体步骤如下：

1. 定义分组窗口：按quantity、category、activity分组，按participants降序排序，用于标记每组内的行优先级。
2. 添加排名标记：在分组窗口内计算每行的排名，排名为1的是该组内participants最大的行。
3. 修改cost字段：使用when条件判断，若行的排名大于1（即participants更小），则将cost设为0，否则保留原cost值。

完整代码示例：

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

# 定义窗口：按目标组合分组，按participants降序排序
window_spec = Window.partitionBy('quantity', 'category', 'activity').orderBy(f.desc('participants'))

# 计算排名并修改cost列
result_df = deptDF.withColumn(
    'rank', f.rank().over(window_spec)
).withColumn(
    'cost',
    f.when(f.col('rank') > 1, 0).otherwise(f.col('cost'))
).drop('rank')  # 移除临时排名列

result_df.show()

执行结果 #

处理后的DataFrame如下：

+--------+----+------------+----------+--------+-----------+
|quantity|cost|participants|activity  |category|id         |
+--------+----+------------+----------+--------+-----------+
|4       |10  |2           |skiing    |outdoor |8589934592 |
|4       |0   |3           |golf      |indoor  |17179869184|
|4       |10  |5           |swimming  |outdoor |25769803776|
|4       |10  |3           |basketball|indoor  |34359738368|
|4       |11  |7           |pool      |indoor  |42949672960|
|4       |11  |12          |pool      |outdoor |51539607552|
|4       |13  |15          |golf      |indoor  |60129542144|
+--------+----+------------+----------+--------+-----------+

方案优势 #

• 全程基于分布式DataFrame操作，避免了collect()带来的性能问题，适合大数据场景。
• 代码简洁，逻辑清晰，无需额外的迭代操作。

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