如何添加具有特殊条件的不同行的两列？

Question

嗨，我有一个 PySpark 数据框。所以，我想从具有特殊条件的不同行中添加两列。其中一列是日期类型。

这里是数据的例子：

--------------------------------
| flag|      date     |  diff  |
--------------------------------
| 1   |   2014-05-31  | 0      |
--------------------------------
| 2   |   2014-06-02  | 2      |
--------------------------------
| 3   |   2016-01-14  | 591    |
--------------------------------
| 1   |   2016-07-08  | 0      |
--------------------------------
| 2   |   2016-07-12  | 4      |
--------------------------------

目前我只知道如何添加两列，使用此代码：

from pyspark.sql.functions import expr
dataframe.withColumn("new_column", expr("date_add(date_column, int_column)"))

预期结果：

有一个名为 "new_date" 的新列，它是将 "diff" 列添加到 "date column" 的结果。

要注意的是有一个特殊条件：如果 "flag" 为 1，则 "date" 和 "diff" 来自同一行，如果不是，则 "date"来自上一行.

我知道在这种情况下，我的数据必须正确排序。

如果有人能帮助我，我将不胜感激。谢谢。

Answer 1

您只需使用 Window 创建一个包含前一个日期的列，并根据 'flag'

的值构建新列

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

w = Window().partitionBy().orderBy(F.col('date'))

dataframe = dataframe.withColumn('previous_date', F.lag('date', 1).over(w))

dataframe = dataframe.withColumn('new_date',
                                 F.when(F.col('flag')==1,
                                        F.expr("date_add(previous_date, diff)")
                                        ).otherwise(F.expr("date_add(date, diff)"))
                                ).drop('previous_date')

Answer 2

以防万一您对 Xavier 的回答有同样的问题。这个想法是一样的，但是我删除了 Window 的一些不必要的条件并修复了语法错误，以及我在尝试他的版本时遇到的 date_add 错误。

from pyspark.sql.functions import *
df1 = spark.createDataFrame([(1,datetime.date(2014,5,31),0),(2,datetime.date(2014,6,2),2),(3,datetime.date(2016,1,14),591),(1,datetime.date(2016,7,8),0),(2,datetime.date(2016,7,12),4)], ["flag","date","diff"])

w = Window.orderBy(col("date"))
df1 = df1.withColumn('previous_date', lag('date', 1).over(w))
df1 = df1.withColumn('new_date',when(col('flag')==1,\
expr('date_add(date, diff)'))\
.otherwise(expr('date_add(previous_date,diff)'))).drop('previous_date')
df1.show()

输出：

+----+----------+----+----------+
|flag|      date|diff|  new_date|
+----+----------+----+----------+
|   1|2014-05-31|   0|2014-05-31|
|   2|2014-06-02|   2|2014-06-02|
|   3|2016-01-14| 591|2016-01-14|
|   1|2016-07-08|   0|2016-07-08|
|   2|2016-07-12|   4|2016-07-12|
+----+----------+----+----------+

如何添加具有特殊条件的不同行的两列？

How do I add two columns with different rows with special condition?

python

dataframe

apache-spark

pyspark

databricks