python - 获取列的“圆形滞后”

我想基于现有列的滞后值在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|
+---+-----+----------------+-----------------+

我可以感觉到它与窗口函数或pyspark.sql.lag函数有关，但不知道该怎么做。

最佳答案

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

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|
+---+-----+----------------+-----------------+

关于python - 获取列的“圆形滞后”，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/51731342/