Python 如何在pyspark中根据多个条件选择前n行项目

现在我有这样的数据：

+----+----+
|col1|   d|
+----+----+
|   A|   4|
|   A|  10|
|   A|   3|
|   B|   3|
|   B|   6|
|   B|   4|
|   B| 5.5|
|   B|  13|
+----+----+

col1是StringType，d是TimestampType，这里我使用DoubleType。 我想根据条件元组生成数据。 给定一个元组[（a，3.5），（a，8），（B，3.5），（B，10）] 我想要这样的结果

+----+---+
|col1|  d|
+----+---+
|   A|  4|
|   A| 10|
|   B|  4|
|   B| 13|
+----+---+

也就是说，对于元组中的每个元素，我们从pyspark数据帧中选择d大于元组数且col1等于元组字符串的第一行。 我已经写的是：

df_res=spark_empty_dataframe    
for (x,y) in tuples:
         dft=df.filter(df.col1==x).filter(df.d>y).limit(1)
         df_res=df_res.union(dft)

---

但我认为这可能存在效率问题，我不知道我是否正确。

避免循环的一种可能方法是从作为输入的元组创建数据帧：

t = [('A',3.5),('A',8),('B',3.5),('B',10)]
ref=spark.createDataFrame([(i[0],float(i[1])) for i in t],("col1_y","d_y"))

然后我们可以根据条件连接输入数据帧（

df

），然后对元组的键和值进行分组，这些键和值将被重复，以获得每个组的第一个值，然后删除额外的列：

(df.join(ref,(df.col1==ref.col1_y)&(df.d>ref.d_y),how='inner').orderBy("col1","d")

.groupBy("col1_y","d_y").agg(F.first("col1").alias("col1"),F.first("d").alias("d"))

.drop("col1_y","d_y")).show()

---

注意，如果数据帧的顺序很重要，您可以使用

单调递增的\u id

分配索引列，并将它们包括在聚合中，然后按索引列排序

用另一种方式编辑而不是排序，然后直接用

min

获取

first

：

(df.join(ref,(df.col1==ref.col1_y)&(df.d>ref.d_y),how='inner')

.groupBy("col1_y","d_y").agg(F.min("col1").alias("col1"),F.min("d").alias("d"))

.drop("col1_y","d_y")).show()

+----+----+
|col1|   d|
+----+----+
|   B| 4.0|
|   B|13.0|
|   A| 4.0|
|   A|10.0|
+----+----+

---

@anky很抱歉，数据把你弄糊涂了。我已经编辑了我的问题单，数据过滤只是关于d和col1，与其他列无关。谢谢，真的帮了我很多忙。

(df.join(ref,(df.col1==ref.col1_y)&(df.d>ref.d_y),how='inner')

.groupBy("col1_y","d_y").agg(F.min("col1").alias("col1"),F.min("d").alias("d"))

.drop("col1_y","d_y")).show()

+----+----+
|col1|   d|
+----+----+
|   B| 4.0|
|   B|13.0|
|   A| 4.0|
|   A|10.0|
+----+----+