如何用Scala为DataFrame添加重复计数列？现有方法结果不符

我明白你的问题啦——你想要给原始DataFrame的每一行都加上对应Id的重复次数，但用groupBy($"Id").count()得到的是聚合后的汇总结果（每个Id只显示一行和计数），而不是保留所有原始行并附加计数，这就是结果不符合预期的原因。

要实现你的需求，我们需要用到窗口函数（Window Functions），它可以在不聚合行的前提下，对分组内的行进行计算。

解决方案（Scala版本） #

假设你用的是Scala语言的Spark API，代码如下：

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.count

// 定义窗口：按Id分区，将相同Id的行归为一组
val idWindow = Window.partitionBy("Id")

// 给原始DF添加repeatedcount列，计算每个Id分组的总行数
val resultDF = originalDF.withColumn(
  "repeatedcount",
  count("Id").over(idWindow)
)

// 展示最终结果
resultDF.show()

代码逻辑解释 #

• 首先导入窗口函数相关的依赖：Window用于定义分组规则，count用于统计分组内的行数。
• Window.partitionBy("Id")指定了分组依据为Id，所有相同Id的行会被划入同一个计算窗口。
• count("Id").over(idWindow)会在每个窗口内统计总行数，然后把这个数值附加到该窗口的每一行上，这样就保留了原始的所有行数据，同时添加了对应的重复计数。

执行结果 #

运行后你会得到和预期几乎一致的输出（注：你的期望输出里Id=3的repeatedcount写为3是笔误，实际该Id仅出现1次，代码会正确计算为1）：

+---+----+-------------+
| Id|Name|repeatedcount|
+---+----+-------------+
|  1|  ab|            2|
|  1|  Cd|            2|
|  2|  BC|            2|
|  2|  mf|            2|
|  3|  Hh|            1|
+---+----+-------------+

补充：PySpark版本 #

如果你使用的是Python语言的PySpark，逻辑完全一致，仅语法略有差异：

from pyspark.sql import Window
from pyspark.sql.functions import count

id_window = Window.partitionBy("Id")
result_df = original_df.withColumn(
  "repeatedcount",
  count("Id").over(id_window)
)

result_df.show()

内容的提问来源于stack exchange，提问作者Ravikumar Reddy Yeruva