有条件的Pyspark累积总和

时间:2018-07-11 18:31:39

标签: dataframe pyspark apache-spark-sql cumulative-sum

我有一个3列的pyspark数据框:

ID ,每次出现多次; DATE ; DELAY ,如果按时付款则为0,否则为1。

已经按 ID DATE 进行了排序。 我需要创建一个名为 CONSECUTIVE 的列,该列显示为每个ID用 DELAY = 1 连续支付了多少张连续账单。

数据示例和预期结果:

ID    | DATE  | DELAY  | CONSECUTIVE
101   | 1     | 1      | 1
101   | 2     | 1      | 2
101   | 3     | 1      | 3
101   | 4     | 0      | 0
101   | 5     | 1      | 1
101   | 6     | 1      | 2
213   | 1     | 1      | 1
213   | 2     | 1      | 2

有没有不使用熊猫的方法吗?如果是这样,我该怎么办?

1 个答案:

答案 0 :(得分:0)

您可以在window的帮助下进行3转换。

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

df = sqlContext.createDataFrame([
    (101, 1, 1),
    (101, 2, 1), # dasd
    (101, 3, 0),
    (101, 4, 1)
], ["id", 'date', 'delay'])

window = Window.partitionBy('id').orderBy('date')
last_value = F.last('rank').over(window.rowsBetween(-2, -1))
consecutive = F.when( F.col('delay')==0, 0) \
            .otherwise( F.when(F.col('last_rank').isNull(), 1) \
            .otherwise( F.col('last_rank')+1))

df \
    .withColumn('rank', F.row_number().over(window)) \
    .withColumn('rank', F.when(F.col('delay')!=0, F.col('rank')).otherwise(0)) \
    .withColumn('last_rank', last_value) \
    .withColumn('consecutive', consecutive).show()

结果:

+---+----+-----+----+---------+-----------+
| id|date|delay|rank|last_rank|consecutive|
+---+----+-----+----+---------+-----------+
|101|   1|    1|   1|     null|          1|
|101|   1|    1|   2|        1|          2|
|101|   1|    0|   0|        2|          0|
|101|   1|    1|   4|        0|          1|
+---+----+-----+----+---------+-----------+