Pyspark -根据其他列值移位列值

from pyspark.sql import SparkSession
from pyspark.sql.functions import lag
from pyspark.sql.window import Window

# 创建SparkSession
spark = SparkSession.builder.getOrCreate()

# 创建示例数据集
data = [(1, 'A', 10), (2, 'B', 20), (3, 'C', 30), (4, 'D', 40)]
df = spark.createDataFrame(data, ['id', 'col1', 'col2'])

# 定义窗口规范
windowSpec = Window.orderBy('id')

# 使用lag函数获取col2列的移位值，并将结果保存到新列shifted_col2中
df = df.withColumn('shifted_col2', lag('col2').over(windowSpec))

# 显示结果
df.show()

+---+----+----+-------------+
| id|col1|col2|shifted_col2 |
+---+----+----+-------------+
|  1|   A|  10|         null|
|  2|   B|  20|           10|
|  3|   C|  30|           20|
|  4|   D|  40|           30|
+---+----+----+-------------+