我试过Window()，但无法得到多个值。所以，使用UDF，你可以像下面这样实现你的要求。

样本起始数据：

首先，使用group by并将row_id的排序数组作为一个新列并将其存储在一个新的嵌套框架中。

from pyspark.sql.types import *
from pyspark.sql.functions import col, sort_array,collect_list

#Create another dataframe with required row_id list
df2=df.groupBy("Team Name").agg(collect_list("row_id").alias("new_col")).withColumn("new_col",sort_array("new_col"))

#Join it with our dataframe
Joined_df = df.join(df2,  'Team Name',  'inner')
Joined_df.show()

它将给予像下面这样的框架。

现在，为case条件创建一个UDF，并像下面这样从框架中调用它。

# Function to select values based on 'difference' returns the required array
@udf(ArrayType(IntegerType()))
def get_values(arr, diff):
    if diff > 0:
        return arr[-diff:]
    elif diff < 0:
        return arr[:abs(diff)]
    return arr

# Get the required row_id array using the udf
Joined_df.withColumn("new_col", get_values(Joined_df["new_col"], Joined_df["difference"])).show()

结果：