pyspark - 如何在 `groupBy()` 之后选择 DataFrame 的特定行？

Question

我有一个 pyspark DataFrame: df。例如：

 e | attempt | grade
---------------------
 1 | 1       | 100
 2 | 1       | 95
 2 | 2       | 55
 3 | 1       | 78
 3 | 2       | 100
 3 | 3       | 88
.
.
. 

我只想为“e”的每个值选择一行：在具有相同“e”值的所有其他行中，具有最大“尝试”值并且它们的尝试小于 X 的行。

例如，如果我打电话，get_results(3)我应该得到以下信息：

 e | attempt | grade
---------------------
 1 | 1       | 100
 2 | 2       | 55
 3 | 3       | 88
.
.
. 

如果我打电话get_results(2)，我应该得到以下信息：

 e | attempt | grade
---------------------
 1 | 1       | 100
 2 | 2       | 55
 3 | 2       | 100
.
.
. 

如果我打电话get_results(1)，我应该得到以下信息：

 e | attempt | grade
---------------------
 1 | 1       | 100
 2 | 1       | 95
 3 | 1       | 78
.
.
. 

我想我应该从 开始df.groupby('e')，但我不知道如何从那里继续。

标签： pysparkpyspark-dataframes

Answer 1

这个想法是首先使用列e和attempt（以降序）顺序对 DataFrame 进行排序。完成后，我们选择顶行。

# Loading the requisite packages and creating the DataFrame.
from pyspark.sql.window import Window
from pyspark.sql.functions import col, first, row_number

valuesCol = [(1,1,100),(2,1,95),(2,2,55),(3,1,78),(3,2,100),(3,3,88)]
df = spark.createDataFrame(valuesCol,['e','attempt','grade'])
df.show()
+---+-------+-----+
|  e|attempt|grade|
+---+-------+-----+
|  1|      1|  100|
|  2|      1|   95|
|  2|      2|   55|
|  3|      1|   78|
|  3|      2|  100|
|  3|      3|   88|
+---+-------+-----+

现在，我们选择 的值X。正如 OP 所说， 的值attempt不能大于X，所以我们过滤掉所有attempt大于的行，X然后使用orderBy()函数排序。

X=2
w = Window.partitionBy(col('e')).orderBy(col('attempt').desc())    
df = df.where(col('attempt')<=X).orderBy(['e','attempt'], ascending=[1,0])
df.show()
+---+-------+-----+
|  e|attempt|grade|
+---+-------+-----+
|  1|      1|  100|
|  2|      2|   55|
|  2|      1|   95|
|  3|      2|  100|
|  3|      1|   78|
+---+-------+-----+

完成后，我们使用row_number()andwindow函数来选择这个排序后的 DataFrame 中的第一行。

df = df.withColumn('row_num', row_number().over(w)).where(col('row_num') == 1).drop('row_num')
df.show()
+---+-------+-----+
|  e|attempt|grade|
+---+-------+-----+
|  1|      1|  100|
|  3|      2|  100|
|  2|      2|   55|
+---+-------+-----+