PySpark如何追踪变量按列值增减 计算带上下限的累计值生成新列

注：该问题已有细节更详实、解释更清晰的版本并获得可行解决方案。

需求说明 #

现有DataFrame中存在一列取值仅为1、0、-1，需要新增计算列，规则如下：

• 计算初始值为3
• 列的取值严格限制在1到5的闭区间内
• 逐行计算时，以上一行的最终计算结果为基数，加上当前行1/0/-1列的取值，若计算结果超出1-5的范围，则直接截断到最近的边界值

期望计算结果样例：

+---+-----+
|var|expec|
+---+-----+
|  1|  4  |
|  1|  5  |
|  0|  5  |
| -1|  4  |
|  1|  5  |
|  1|  5  |
| -1|  4  |
+---+-----+

已尝试的实现思路 #

• 窗口函数实现：尝试使用窗口函数做累计求和，但lag函数不支持动态取值，无法满足递归依赖上一行截断后结果的要求——普通窗口累计和是先算完所有总和再截断，和逐行依赖上一步截断结果的逻辑不匹配。
• reduce逐行迭代：尝试通过reduce方法逐行迭代累计计算值，但始终没有理清reduce函数中累计行结果的返回逻辑，已编写的未完成代码如下：

from pyspark.sql import Row

def calc_rec(acm, row):
    acm_ = []
    if (type(acm) == type([])):
        cur = acm[1][-1]
        res = acm[0]
        rows_ = acm[1]
    elif (acm == 3):
        res = acm
        cur = row
        rows_ = []
        rows_.append(Row(resc=acm,valor=row.valor))
    x = res + cur.valor
    if x > 5:
        x = 5
    elif x <= 0:
        x = 1
    rows_.append(Row(resc=x, valor=cur.valor))
    acm_ = [x, rows_]
    return acm_

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