PySpark中按分区生成递增Group_number列的实现方案

核心逻辑 #

按registerNumber和share_class_number分组，组内遵循以下规则：

• 首次出现非零no_of_shares时，Group_number赋值为1
• 每遇到下一个非零no_of_shares，Group_number递增1
• 零值行的Group_number继承最近一次非零行的数值
• 切换分组（registerNumber或share_class_number变更）时，Group_number重置为1

实现步骤 #

1. 导入依赖库

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

2. 定义窗口分区与排序规则
   窗口按主键registerNumber和share_class_number分区，组内按行顺序排序（若有业务专属排序字段，可替换monotonically_increasing_id()）：

window_spec = Window.partitionBy("registerNumber", "share_class_number")\
                    .orderBy(F.monotonically_increasing_id())

3. 生成Group_number列
   先标记非零行，再通过累计求和实现分组内的递增计数：

# 标记非零行
df = df.withColumn("is_non_zero", F.when(F.col("no_of_shares") != 0, 1).otherwise(0))

# 计算累计和得到Group_number
df_result = df.withColumn(
    "Group_number",
    F.sum("is_non_zero").over(window_spec)
).drop("is_non_zero")

示例验证 #

假设输入DataFrame：

处理后输出：

特殊情况处理 #

如果需要将组内未出现非零值的零行Group_number设为1（而非0），可添加额外逻辑：

df_result = df_result.withColumn(
    "Group_number",
    F.when(F.col("Group_number") == 0, 1).otherwise(F.col("Group_number"))
)

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