PySpark DataFrame中id_sum列0值替换为1到n序列的代码问题排查

原始数据与需求 #

原始PySpark DataFrame df_model：

需求是为每个id_client分组内的行生成从1开始的递增计数（重复行也需依次增加），期望输出：

现有代码的问题 #

你尝试的代码（包括调整orderBy为id_sku的版本）存在两个核心问题：

1. 逻辑方向错误：用f.sum(f.when(f.col('id_client') != t, 1).otherwise(0))计算，其中t是第一行的id_client（1111）。这导致id_client=1111的所有行都满足id_client == t，累加值为0，完全不符合分组内计数的需求；其他分组的结果看似部分正确，只是巧合，逻辑本质是统计“非第一个client的行数”，和需求无关。
2. 窗口函数误用：你的需求是分组内行号计数，却用了求和逻辑，偏离了目标。另外，collect()[0][0]的方式不推荐，会将数据拉到Driver端，大数据场景下有性能风险，且完全没必要。

正确实现代码 #

直接使用row_number()窗口函数，按id_client分区，为保证重复行的计数顺序稳定，可结合monotonically_increasing_id()作为排序辅助列：

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

# 定义窗口：按id_client分区，按id_sku+唯一ID排序保证顺序稳定
window_spec = Window.partitionBy("id_client").orderBy("id_sku", f.monotonically_increasing_id())

# 生成id_sum列
df_result = df_model.withColumn("id_sum", f.row_number().over(window_spec))

执行后即可得到你期望的输出：每个id_client分组内的行从1开始依次递增计数，重复行也会按顺序生成连续数字。

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