获取列的“循环滞后”答案

【问题标题】：Get "circular lag" of a column获取列的“循环滞后”
【发布时间】：2019-01-14 19:30:52
【问题描述】：

我想根据现有列的滞后值在 pyspark.sql.DataFrame 中创建一个新列。但是......我也希望最后一个值成为第一个值，第一个值成为最后一个值。这是一个例子：

df = spark.createDataFrame([(1,100),
                            (2,200),
                            (3,300),
                            (4,400),
                            (5,500)],
                            ['id','value'])

df.show()

+---+-----+
| id|value|
+---+-----+
|  1|  100|
|  2|  200|
|  3|  300|
|  4|  400|    
|  5|  500|
+---+-----+

期望的输出是：

+---+-----+----------------+-----------------+
| id|value|lag_value_plus_2|lag_value_minus_2|
+---+-----+----------------+-----------------+
|  1|  100|             300|              400|
|  2|  200|             400|              500|
|  3|  300|             500|              100|
|  4|  400|             100|              200|
|  5|  500|             200|              300|
+---+-----+----------------+-----------------+

感觉和window函数或者pyspark.sql.lag函数有关系，但是不知道怎么做。

【问题讨论】：

如果您知道每个 Window 分区的元素数量（在这种情况下为 5，因为我假设 Window 只是按 id 排序），我可以想办法做到这一点。
这似乎是可行的，因为如果需要，我们可以在数据框中添加一列，其中包含每个 Window 分区的元素数，使用如下：df.withColumn('n_elem',F.count().over(w))?

标签： python pyspark pyspark-sql window-functions

【解决方案1】：

这是我可以提供的一种解决方案。但我不确定它是否是最优化的：

from functools import reduce                                                                                       

# Duplicate the dataframe twice, one "before" and one "after"
df = reduce(
    lambda a, b : a.union(b), 
    [df.withColumn("x", F.lit(i)) for i in [-1,0,1]] 
)

df.withColumn(
    "lag_value_plus_2",
    F.lead("value", 2).over(Window.partitionBy().orderBy("x", "id"))
).withColumn(
    "lag_value_minus_2",
    F.lag("value", 2).over(Window.partitionBy().orderBy("x", "id"))
).where("x=0").drop("x").show()

+---+-----+----------------+-----------------+
| id|value|lag_value_plus_2|lag_value_minus_2|
+---+-----+----------------+-----------------+
|  1|  100|             300|              400|
|  2|  200|             400|              500|
|  3|  300|             500|              100|
|  4|  400|             100|              200|
|  5|  500|             200|              300|
+---+-----+----------------+-----------------+

【讨论】：