如何为每个用户（另一列）仅保留列值最高的pyspark dataframe行？

Question

我有一个Pypark数据帧-

df = spark.createDataFrame([
    ("u1", "india", 10),
    ("u1", "australia", 12),
    ("u2", "usa", 14),
    ("u2", "nepal", 7),
    ("u3", "england", 13),
    ("u4", "japan", 15),],
    ['user_id', 'place', 'price'])

df.show()

输出-

+-------+---------+-----+
|user_id|    place|price|
+-------+---------+-----+
|     u1|    india|   10|
|     u1|australia|   12|
|     u2|      usa|   14|
|     u2|    nepal|    7|
|     u3|  england|   13|
|     u4|    japan|   15|
+-------+---------+-----+

我只想为每个用户id

所需的输出应如下所示：

+-------+---------+-----+
|user_id|    place|price|
+-------+---------+-----+
|     u1|australia|   12|
|     u2|      usa|   14|
|     u3|  england|   13|
|     u4|    japan|   15|
+-------+---------+-----+

我如何才能做到这一点

Answer 1

您可以使用group by和加入表格来获得位置。根据我的经验，Groupby的性能优于Windows：

import pyspark.sql.functions as f

df2 = df.groupby("user_id").agg(f.max("price").alias("price"))
df2 = df2.join(df.dropDuplicates(), ["user_id", "price"])

df2.show()

Answer 2

为每个分区中的每一行分配一个行号，并过滤行号为1的行

import pyspark.sql.functions as F
from pyspark.sql.window import Window

df2 = df.withColumn(
    'rn',
    F.row_number().over(Window.partitionBy('user_id').orderBy(F.desc('price')))
).filter('rn = 1').drop('rn').orderBy('user_id')

df2.show()
+   -+    -+  -+
|user_id|    place|price|
+   -+    -+  -+
|     u1|australia|   12|
|     u2|      usa|   14|
|     u3|  england|   13|
|     u4|    japan|   15|
+   -+    -+  -+