如何 select 基於多個條件的前 n 行項目 pyspark

Question

現在我有這樣的數據：

+----+----+
|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 dataframe 中的 select dataframe 的前 1 行 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)

但我認為這可能有效率問題，我不知道我是否正確。

Answer 1

避免循環的一種可能方法是從您作為輸入的元組創建 dataframe ：

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()

---

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

請注意，如果 dataframe 的順序很重要，您可以使用monotonically_increasing_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|
+----+----+