PySpark-将上一行和下一行追加到当前行

时间:2018-07-10 15:03:06

标签: python apache-spark dataframe pyspark apache-spark-sql

假设我有一个像这样的PySpark数据框:

list(set(df['time'].dt.year.values))

如何将一行的最后一列和下一列追加到当前行,如下所示:

1 0 1 0
0 0 1 1
0 1 0 1

我熟悉用于添加列的1 0 1 0 0 0 0 0 0 0 1 1 0 0 1 1 1 0 1 0 0 1 0 1 0 1 0 1 0 0 1 1 0 0 0 0 方法,但是不确定在该字段中输入什么内容。

.withColumn()是占位符值,因为在这些行之前和之后没有之前或之后的行。

1 个答案:

答案 0 :(得分:2)

您可以使用pyspark.sql.functions.lead()pyspark.sql.functions.lag(),但首先需要一种排序列的方法。如果您还没有确定顺序的列,则可以使用pyspark.sql.functions.monotonically_increasing_id()

然后将其与Window函数结合使用。

例如,如果您具有以下DataFrame df

df.show()
#+---+---+---+---+
#|  a|  b|  c|  d|
#+---+---+---+---+
#|  1|  0|  1|  0|
#|  0|  0|  1|  1|
#|  0|  1|  0|  1|
#+---+---+---+---+

您可以这样做:

from pyspark.sql import Window
import pyspark.sql.functions as f

cols = df.columns
df = df.withColumn("id", f.monotonically_increasing_id())
df.select(
    "*", 
    *([f.lag(f.col(c),default=0).over(Window.orderBy("id")).alias("prev_"+c) for c in cols] + 
      [f.lead(f.col(c),default=0).over(Window.orderBy("id")).alias("next_"+c) for c in cols])
).drop("id").show()
#+---+---+---+---+------+------+------+------+------+------+------+------+
#|  a|  b|  c|  d|prev_a|prev_b|prev_c|prev_d|next_a|next_b|next_c|next_d|
#+---+---+---+---+------+------+------+------+------+------+------+------+
#|  1|  0|  1|  0|     0|     0|     0|     0|     0|     0|     1|     1|
#|  0|  0|  1|  1|     1|     0|     1|     0|     0|     1|     0|     1|
#|  0|  1|  0|  1|     0|     0|     1|     1|     0|     0|     0|     0|
#+---+---+---+---+------+------+------+------+------+------+------+------+