如何在PySpark中同时使用groupby与nth_value实现分组取指定行效果

PySpark 没有提供 Pandas 中groupby.nth的直接API，可通过窗口函数加行号筛选的方式实现对应逻辑，该方案适配大数据量生产场景，无内存溢出风险，实现步骤如下：

核心逻辑说明 #

因为Spark是分布式计算引擎，默认不会保留原始数据的读取顺序，所以需要先给全表数据加全局唯一递增ID作为排序依据，保证分组内的行顺序和你原始数据集顺序一致，再按日期分组加组内行号，筛选对应位置的行即可。

完整实现代码 #

# 导入依赖
from pyspark.sql import Window
import pyspark.sql.functions as F

# --------------------------
# 第一步：构造示例数据集（你实际使用时替换成自己的df_apps即可）
# --------------------------
data = [
    ("2021-01-10", "FACEBOOK", 1000, 5000),
    ("2021-01-10", "FACEBOOK", 20000, 900000),
    ("2021-02-10", "FACEBOOK", 9000, 72000),
    ("2021-01-11", "FACEBOOK", 4000, 2000),
    ("2021-01-11", "FACEBOOK", 40000, 85000),
    ("2021-02-11", "FACEBOOK", 1000, 2000)
]
schema = ["DATE", "APP", "DOWNLOADS", "ACTIVE_USERS"]
df_apps = spark.createDataFrame(data, schema=schema)

# --------------------------
# 第二步：实现分组取第n行逻辑
# --------------------------
# 加全局行号，保证分组内排序和原始数据顺序一致
df_with_id = df_apps.withColumn("global_row_id", F.monotonically_increasing_id())

# 定义窗口：按DATE分组，按全局行号升序排序
window_spec = Window.partitionBy("DATE").orderBy("global_row_id")

# 给每个分组内的行加1-based的行号
df_with_row_num = df_with_id.withColumn("group_row_num", F.row_number().over(window_spec))

# 筛选组内行号为2的行（对应Pandas nth(1)的0-based索引），删除辅助列得到最终结果
df_result = df_with_row_num.filter(F.col("group_row_num") == 2).drop("global_row_id", "group_row_num")

# 查看结果
df_result.show()

输出结果 #

+----------+--------+---------+------------+
|      DATE|     APP|DOWNLOADS|ACTIVE_USERS|
+----------+--------+---------+------------+
|2021-01-10|FACEBOOK|    20000|      900000|
|2021-01-11|FACEBOOK|    40000|       85000|
+----------+--------+---------+------------+

自定义调整说明 #

• 如果你需要按指定字段排序后取中间行，只需修改窗口定义里的orderBy参数即可，比如按下载量降序排序：window_spec = Window.partitionBy("DATE").orderBy(F.desc("DOWNLOADS"))
• 如果你要取的位置不是0-based的第1位，修改筛选条件里的数值即可，比如要取0-based第2位，就把group_row_num == 2改成group_row_num == 3

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