如何使用pyspark设置动态where子句

我有一个数据集，其中有多个组。我有一个排名列，它递增地统计每个组的每个条目。该结构的示例如下所示：

+-----------+---------+---------+
|  equipment|   run_id|run_order|
+-----------+---------+---------+
|1          |430032589|        1|
|1          |430332632|        2|
|1          |430563033|        3|
|1          |430785715|        4|
|1          |431368577|        5|
|1          |431672148|        6|
|2          |435497596|        1|
|1          |435522469|        7|

每组（设备）有不同的运行量。如上所示，设备1运行7次，而设备2运行1次。我想选择每个设备的第一次和最后n次运行。要选择前n次运行非常简单：

df.select("equipment", "run_id").distinct().where(df.run_order <= n).orderBy("equipment").show()

我可以运行groupBy以获得每个设备的最高运行订单

+-----------+----------------+
|  equipment| max(run_order) |
+-----------+----------------+
|1          |               7|
|2          |               1|

---

但是我不确定是否有一种方法可以构造一个动态where子句，它可以这样工作。这样我就可以得到最后n次运行（包括每次运行的所有时间步数据）。

您可以为每个设备添加一列最大秩，并根据该列进行筛选：

from pyspark.sql import functions as F, Window

n = 3

df2 = df.withColumn(
    'max_run', 
    F.max('run_order').over(Window.partitionBy('equipment'))
).where(F.col('run_order') >= F.col('max_run') - n)

from pyspark.sql import functions as F, Window

n = 3

df2 = df.withColumn(
    'max_run', 
    F.max('run_order').over(Window.partitionBy('equipment'))
).where(F.col('run_order') >= F.col('max_run') - n)