根据得分pyspark筛选其他列中每个值的列前N个值

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, row_number
from pyspark.sql.window import Window

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

# 创建示例数据
data = [
    (1, 10, 20, 30),
    (2, 15, 25, 35),
    (3, 5, 15, 25),
    (4, 20, 30, 40)
]

# 创建DataFrame
df = spark.createDataFrame(data, ["score", "column1", "column2", "column3"])
df.show()

# 定义窗口规范
windowSpec = Window.orderBy(col("score").desc())

# 使用row_number()函数为每个值的列排序
df = df.withColumn("rank", row_number().over(windowSpec))

# 按照得分筛选出前N个值
N = 2
df_filtered = df.filter(col("rank") <= N)

df_filtered.show()