如何在 `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" 的所有其他行的最大 "attempt" 值并且它们的尝试小于 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') 开始，但我不知道如何从那里继续。

Answer 1

我们的想法是首先使用列 e 和 attempt（降序）对 DataFrame 进行排序。完成后，我们 select 顶行。

# 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|
+---+-------+-----+

现在，我们 select 的值为 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() and window 函数来 select 这个已排序 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|
+---+-------+-----+

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

How can I choose a particular line of a DataFrame after `groupBy()`?

pyspark

pyspark-dataframes