如何在PySpark DataFrame中按指定列统计重复行数？

需求说明 #

统计具有相同X、Y和DATE列值的行数，将统计结果存入新列SUM中，同时保留原DataFrame的所有行。

解决方案 #

使用PySpark的窗口函数实现：先按指定列分组，计算每组的总行数，再将统计值映射到组内的每一行。

代码实现 #

1. 初始化环境并创建示例DataFrame

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

spark = SparkSession.builder.appName("group_count_demo").getOrCreate()

# 示例数据
data = [
    (23, 41, 63, "2016-01-01"),
    (23, 41, 5, "2016-01-01"),
    (23, 41, 75, "2016-01-01"),
    (23, 41, 46, "2016-12-01"),
    (23, 41, 23, "2016-12-01"),
    (27, 41, 5, "2016-01-01"),
    (27, 41, 75, "2016-01-01"),
    (27, 41, 85, "2016-01-01"),
    (27, 41, 71, "2016-01-01")
]

df = spark.createDataFrame(data, ["X", "Y", "Z", "DATE"])

2. 定义窗口分组规则

# 按X、Y、DATE列分组的窗口规范
window_spec = Window.partitionBy("X", "Y", "DATE")

3. 添加统计列

# 统计每组行数，生成SUM列
result_df = df.withColumn("SUM", count("*").over(window_spec))

4. 查看结果

result_df.show()

输出结果 #

+---+---+---+----------+---+
|  X|  Y|  Z|      DATE|SUM|
+---+---+---+----------+---+
| 23| 41| 63|2016-01-01|  3|
| 23| 41|  5|2016-01-01|  3|
| 23| 41| 75|2016-01-01|  3|
| 23| 41| 46|2016-12-01|  2|
| 23| 41| 23|2016-12-01|  2|
| 27| 41|  5|2016-01-01|  4|
| 27| 41| 75|2016-01-01|  4|
| 27| 41| 85|2016-01-01|  4|
| 27| 41| 71|2016-01-01|  4|
+---+---+---+----------+---+

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