如何在groupBy（）之后选择数据框的特定行？

2024-05-03 • 问答

我有一个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')开始，但是我不知道如何从那里继续。

这个想法是首先以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,95),2,55),(3,78),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的所有行，然后使用{{ 3}}功能。

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

完成后，我们将使用orderBy()和row_number()函数在此排序的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（）之后选择数据框的特定行？

liupan7891 回答：如何在groupBy（）之后选择数据框的特定行？

大家都在问