如何基于rank列将PySpark DataFrame划分为固定规模分组

需要为PySpark DataFrame新增一列segment，根据已有rank列的值按固定数量（每100,000个rank值为一组）划分分组：

• rank值1~100000 → segment=1
• rank值100001~200000 → segment=2
• 以此类推，需支持任意规模的rank数据。

预期输出示例：

--------------------------------------
| id        |     rank   | segment   | 
--------------------------------------
| 100       |   1        |     1     |  
| 200       |   100002   |     2     | 
| 300       |   900007   |     9     |   
--------------------------------------

解决方案 #

可以通过PySpark的数学函数实现固定数量的分组划分，核心逻辑是对rank值进行整数运算，计算其所属的分组编号。以下提供两种等价的实现方式：

方法1：基于整数除法的偏移计算 #

利用(rank - 1) // 分组大小 + 1的逻辑，避免边界值（如rank=100000）被错误划分到下一组：

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

# 初始化SparkSession
spark = SparkSession.builder.appName("SegmentRank").getOrCreate()

# 创建示例DataFrame
import pandas as pd
df = spark.createDataFrame(pd.DataFrame({
  "id": [100,200,300], "rank": [1, 100002, 900007]
}))

# 定义分组大小
GROUP_SIZE = 100000

# 新增segment列
df_with_segment = df.withColumn(
    "segment",
    F.floor((F.col("rank") - 1) / GROUP_SIZE) + 1
)

# 查看结果
df_with_segment.show()

方法2：基于向上取整函数 #

使用ceil函数直接对rank值除以分组大小的结果取整，逻辑更直观：

df_with_segment = df.withColumn(
    "segment",
    F.ceil(F.col("rank") / GROUP_SIZE)
)

结果验证 #

两种方法都会得到符合预期的输出：

+---+-------+-------+
| id|   rank|segment|
+---+-------+-------+
|100|      1|      1|
|200| 100002|      2|
|300|900007|      9|
+---+-------+-------+

灵活性说明 #

只需修改GROUP_SIZE变量的值，即可适配不同的分组数量要求，完全支持任意规模的rank数据。

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