时间戳

首页 > 解决方案 > 如何在 Pyspark 中编写 for 循环或情节

python - 如何在 Pyspark 中编写 for 循环或情节

问题描述

我有一个数据框，我正在尝试在其上编写一个 for 循环。

|ID   | from_dt  | To_dt    |row_number|diff|negetive_or_not|
+----------+----------+----------+----------+----+---------------+
|11111|2020-07-30|2020-07-31|         1|  -2|              0|
|11111|2020-08-02|2020-08-11|         2|   4|              1|
|11111|2020-08-07|2020-08-08|         3|  -4|              0|
|11111|2020-08-12|2020-08-18|         4|   1|              1|
|11111|2020-08-17|2020-08-19|         5|   0|              1|
|11111|2020-08-19|2020-08-22|         6|   2|              1|
|11111|2020-08-20|2020-08-24|         7|  -1|              0|
|11111|2020-08-25|2020-08-27|         8|   0|              1|
|11111|2020-08-27|2020-08-31|         9|-999|              0|

目标是确定情节。如果负数或不以0开头，则为一集，如果负数或非以1开头，直到达到0，即为一集。

这是理想的输出

+----------+----------+----------+----------+----+---------------+
|ID   | from_dt  | To_dt    |row_number|diff|negetive_or_not| Episode
+----------+----------+----------+----------+----+---------------+
|11111|2020-07-30|2020-07-31|         1|  -2|              0|   1
|11111|2020-08-02|2020-08-11|         2|   4|              1|   2
|11111|2020-08-07|2020-08-08|         3|  -4|              0|   2
|11111|2020-08-12|2020-08-18|         4|   1|              1|   3 
|11111|2020-08-17|2020-08-19|         5|   0|              1|   3 
|11111|2020-08-19|2020-08-22|         6|   2|              1|   3 
|11111|2020-08-20|2020-08-24|         7|  -1|              0|   3 
|11111|2020-08-25|2020-08-27|         8|   0|              1|   4
|11111|2020-08-27|2020-08-31|         9|-999|              0|   4
|22222|2020-07-30|2020-07-31|         1|  -2|              0|   1
|22222|2020-08-02|2020-08-11|         2|   4|              1|   2
|22222|2020-08-07|2020-08-08|         3|  -4|              0|   2
+----------+----------+----------+----------+----+---------------+

我尝试使用 case when 和 rank，例如

当 negetive_or_not = 0 然后是 "eps1" 或者 "eps2" 时，两者都不起作用。

df2 = df.selectExpr('*') .withColumn("Episode",lead(col("to_dt")).over(Window.partitionBy("patient_id").orderBy(col("negetive_or_not"))))

我也尝试在 pyspark 中编写一个 for 循环，但我很难将数据帧传输到列表中，任何建议都将受到高度赞赏。

标签： pythonloopsfor-looppysparkapache-spark-sql

解决方案

该方法如下

  1. 当值应该改变时，首先标记该行。
  2. 基于标志生成一个新值。
  3. 前向填充空值
df.withColumn('flag', when(((col('negetive_or_not')==1) & (lag('negetive_or_not').over(Window.partitionBy('ID').orderBy('row_number'))==0)) | (lag('negetive_or_not').over(Window.partitionBy('ID').orderBy('row_number')).isNull()), lit('change')).otherwise(lit('no'))).\
   withColumn('ep', when(col('flag')=='change',row_number().over(Window.partitionBy('ID','flag').orderBy('row_number')))).\
   withColumn('Episode', last('ep', ignorenulls=True).over(Window.partitionBy('ID').orderBy('row_number'))).\
   drop('flag','ep').show()

+-----+----------+----------+----------+----+---------------+-------+
|   ID|   from_dt|     To_dt|row_number|diff|negetive_or_not|Episode|
+-----+----------+----------+----------+----+---------------+-------+
|11111|2020-07-30|2020-07-31|         1|  -2|              0|      1|
|11111|2020-08-02|2020-08-11|         2|   4|              1|      2|
|11111|2020-08-07|2020-08-08|         3|  -4|              0|      2|
|11111|2020-08-12|2020-08-18|         4|   1|              1|      3|
|11111|2020-08-17|2020-08-19|         5|   0|              1|      3|
|11111|2020-08-19|2020-08-22|         6|   2|              1|      3|
|11111|2020-08-20|2020-08-24|         7|  -1|              0|      3|
|11111|2020-08-25|2020-08-27|         8|   0|              1|      4|
|11111|2020-08-27|2020-08-31|         9|-999|              0|      4|
+-----+----------+----------+----------+----+---------------+-------+

推荐阅读