pyspark 使用每个 id 的最小值过滤数据帧

Question

给定如下表格：

+--+------------------+-----------+
|id|     diagnosis_age|  diagnosis|
+--+------------------+-----------+
| 1|2.1843037179180302| 315.320000|
| 1|  2.80033330216659| 315.320000|
| 1|   2.8222365762732| 315.320000|
| 1|  5.64822705794013| 325.320000|
| 1| 5.686557787521759| 335.320000|
| 2|  5.70572315231258| 315.320000|
| 2| 5.724888517103389| 315.320000|
| 3| 5.744053881894209| 315.320000|
| 3|5.7604813374292005| 315.320000|
| 3|  5.77993740687426| 315.320000|
+--+------------------+-----------+

我试图通过仅考虑每个 id 诊断年龄最小的诊断来减少每个 id 的记录数量。在 SQL 中，您可以将表连接到自身，例如：

SELECT a.id, a.diagnosis_age, a.diagnosis
    FROM tbl1 a
INNER JOIN
(SELECT id, MIN(diagnosis_age) AS min_diagnosis_age
    FROM tbl1
        GROUP BY id) b
ON b.id = a.id
WHERE b.min_diagnosis_age = a.diagnosis_age

如果是 rdd，你可以这样做：

rdd.map(lambda x: (x["id"], [(x["diagnosis_age"], x["diagnosis"])]))\
.reduceByKey(lambda x, y: x + y)\
.map(lambda x: (x[0], [i for i in x[1] if i[0] == min(x[1])[0]]))

您将如何仅使用 spark 数据帧操作来实现相同的目标？如果这是可能的？特别是没有 sql/rdd 操作。

谢谢

Answer 1

您可以将window 与first 函数一起使用，然后将filter 与所有其他函数一起使用。

from pyspark.sql import functions as F
from pyspark.sql.window import Window
w=Window().partitionBy("id").orderBy("diagnosis_age")
df.withColumn("least_age", F.first("diagnosis_age").over(w))\
.filter("diagnosis_age=least_age").drop("least_age").show()

+---+------------------+---------+
| id|     diagnosis_age|diagnosis|
+---+------------------+---------+
|  1|2.1843037179180302|   315.32|
|  3| 5.744053881894209|   315.32|
|  2|  5.70572315231258|   315.32|
+---+------------------+---------+

你也可以这样做没有窗口功能，使用groupBymin和first：

from pyspark.sql import functions as F
df.orderBy("diagnosis_age").groupBy("id")\
.agg(F.min("diagnosis_age").alias("diagnosis_age"), F.first("diagnosis").alias("diagnosis"))\
.show()
+---+------------------+---------+
| id|     diagnosis_age|diagnosis|
+---+------------------+---------+
|  1|2.1843037179180302|   315.32|
|  3| 5.744053881894209|   315.32|
|  2|  5.70572315231258|   315.32|
+---+------------------+---------+

注意，我在groupyBy 之前订购了由diagnosis_age 处理那些所需的诊断值没有出现在第一行的情况强>组。 但是，如果您的数据已经按diagnosis_age 排序，您可以使用上面的代码没有orderBy。