创建一个新列，当值重置为另一列中的pyspark中的1时，增加值。

Question

逻辑列在DataFrame中考虑一个列，如1，2，3，4，1，2，1，2，3，2，1，2，1，2，2，2，2，2，2，2，2，2，2，2，2，2，2，2，2，2，2，2，4，4，5，2，6，7，2。火花柱创建一个新的列，当值重置为1时，预期产出为1，1，1，2，2，3，4，4，5，5，6，7，7。

我对火种有点陌生，如果有人能帮我的话，那对我来说就太棒了。

如下所示编写了逻辑

def sequence(row_num):
    results = [1, ]
    flag = 1
    for col in range(0, len(row_num)-1):
        if row_num[col][0]>=row_num[col+1][0]:
            flag+=1
        results.append(flag)
    return results

但无法通过udf传递列。请帮我这个忙

Answer 1

你的数据：

df  = spark.createDataFrame(
  [
('1','a'),
('2','b'),
('3','c'),
('4','d'),
('1','e'),
('2','f'),
('1','g'),
('1','h'),
('2','i'),
('3','j'),
('1','k'),
('2','l'),
('1','m'),
('1','n'),
('2','o')
  ], ['group','label']
)

+-----+-----+
|group|label|
+-----+-----+
|    1|    a|
|    2|    b|
|    3|    c|
|    4|    d|
|    1|    e|
|    2|    f|
|    1|    g|
|    1|    h|
|    2|    i|
|    3|    j|
|    1|    k|
|    2|    l|
|    1|    m|
|    1|    n|
|    2|    o|
+-----+-----+

您可以创建一个标志，并使用一个窗口函数来计算累积和。不需要使用UDF：

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

w = W.partitionBy().orderBy('label').rowsBetween(Window.unboundedPreceding, 0)

df\
    .withColumn('Flag', F.when(F.col('group') == 1, 1).otherwise(0))\
    .withColumn('Output', F.sum('Flag').over(w))\
    .show()

+-----+-----+----+------+
|group|label|Flag|Output|
+-----+-----+----+------+
|    1|    a|   1|     1|
|    2|    b|   0|     1|
|    3|    c|   0|     1|
|    4|    d|   0|     1|
|    1|    e|   1|     2|
|    2|    f|   0|     2|
|    1|    g|   1|     3|
|    1|    h|   1|     4|
|    2|    i|   0|     4|
|    3|    j|   0|     4|
|    1|    k|   1|     5|
|    2|    l|   0|     5|
|    1|    m|   1|     6|
|    1|    n|   1|     7|
|    2|    o|   0|     7|
+-----+-----+----+------+