python - 获取某列的 "circular lag"

Question

我想根据现有列的滞后值在 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 函数有关，但不知道该怎么做。

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