如何在PySpark中从源DataFrame生成指定NEW_VERSION列结果？

我有如下DataFrame：

+------------+--------------------+-------------+
|PERSON_NBR  |PERSON_VERSION_NBR  |RECRD_TYPE_CD|
+------------+--------------------+-------------+
|  0065321901|                   1|            1|
|  0065321901|                   2|            4|
|  0065321901|                   3|            5|
|  0065321901|                   4|            2|
|  0065321901|                   5|            6|
|  0065321901|                   6|            7|
|  0065321901|                   5|            2|
|  0065321901|                   6|            8|
|  0775123711|                   1|            1|
|  0775123711|                   2|            3|
|  0775123711|                   3|            2|
|  1237251722|                   1|            1|
|  1237251722|                   2|            3|
+------------+--------------------+-------------+

需要新增一列NEW_VERSION，规则为：

• 初始值为1
• 当某行的RECRD_TYPE_CD为2时，同一PERSON_NBR对应的后续所有记录的NEW_VERSION值递增1

期望输出如下：

+------------+--------------------+-------------+----------------+
|PERSON_NBR  |PERSON_VERSION_NBR  |RECRD_TYPE_CD| NEW_VERSION|
+------------+--------------------+-------------+----------------+
|  0065321901|                   1|            1|               1|
|  0065321901|                   2|            4|               1|
|  0065321901|                   3|            5|               1|
|  0065321901|                   4|            2|               1|
|  0065321901|                   5|            6|               2|
|  0065321901|                   6|            7|               2|
|  0065321901|                   7|            2|               2|
|  0065321901|                   8|            8|               3|
|  0775123711|                   1|            1|               1|
|  0775123711|                   2|            3|               1|
|  0775123711|                   3|            2|               1|
|  1237251722|                   1|            1|               1|
|  1237251722|                   2|            3|               1|
+------------+--------------------+-------------+----------------+

可以通过Pandas的分组、移位和累积求和实现需求，代码如下：

import pandas as pd

# 构造原始DataFrame
data = {
    'PERSON_NBR': ['0065321901']*8 + ['0775123711']*3 + ['1237251722']*2,
    'PERSON_VERSION_NBR': [1,2,3,4,5,6,5,6,1,2,3,1,2],
    'RECRD_TYPE_CD': [1,4,5,2,6,7,2,8,1,3,2,1,3]
}
df = pd.DataFrame(data)

# 新增NEW_VERSION列
df['NEW_VERSION'] = 1 + df.groupby('PERSON_NBR')['RECRD_TYPE_CD'].transform(
    lambda s: (s == 2).shift().fillna(0).cumsum()
)

print(df)

代码说明 #

1. 按PERSON_NBR分组，确保每个用户的版本计算独立
2. 对每组的RECRD_TYPE_CD判断是否等于2，生成布尔序列
3. 使用shift(1)将布尔序列向下移位一行，使当前行的取值对应上一行是否为2
4. 用fillna(0)处理移位后的空值（每组第一行移位后为NaN）
5. 对移位后的序列做累积求和，得到截至当前行之前出现过的RECRD_TYPE_CD=2的次数
6. 最后加1，得到初始值为1、符合规则的NEW_VERSION

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