使用窗口函数的pyspark

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import rank, dense_rank, row_number, lag, lead, sum, avg, min, max

spark = SparkSession.builder.appName("WindowFunctionExample").getOrCreate()

data = [(1, "A", 100), (2, "B", 200), (3, "A", 150), (4, "B", 300), (5, "A", 200)]
df = spark.createDataFrame(data, ["id", "category", "value"])

windowSpec = Window.partitionBy("category").orderBy("value")

df.withColumn("rank", rank().over(windowSpec)) \
  .withColumn("dense_rank", dense_rank().over(windowSpec)) \
  .withColumn("row_number", row_number().over(windowSpec)) \
  .withColumn("lag_value", lag("value").over(windowSpec)) \
  .withColumn("lead_value", lead("value").over(windowSpec)) \
  .withColumn("sum_value", sum("value").over(windowSpec)) \
  .withColumn("avg_value", avg("value").over(windowSpec)) \
  .withColumn("min_value", min("value").over(windowSpec)) \
  .withColumn("max_value", max("value").over(windowSpec)) \
  .show()