groupBy按category_code分组后取count最大值对应brand值的实现方法

要实现按category_code分组取count最大值对应的brand，推荐用窗口函数实现，仅需一次shuffle，大数据量下性能远高于分组后再join的写法，两种常用实现方式如下：

PySpark DataFrame API 写法 #

# 引入依赖函数
from pyspark.sql import functions as F
from pyspark.sql.window import Window

# 定义窗口规则：按category_code分组，组内按count倒序排序
win = Window.partitionBy("category_code").orderBy(F.desc("count"))

# 给组内每行加排名，筛选排名第一的行，就是每组count最大的记录
res_df = df.withColumn("group_rank", F.rank().over(win)) \
           .filter(F.col("group_rank") == 1) \
           .select("category_code", "brand", "count")

# 查看结果
res_df.show()

说明：如果需要同组多个并列最大值时仅返回任意一条，把F.rank()替换为F.row_number()即可。

Spark SQL 写法 #

首先注册临时视图：

df.createOrReplaceTempView("product_data")

执行SQL查询：

select category_code, brand, count
from (
    select 
        category_code,
        brand,
        count,
        rank() over(partition by category_code order by count desc) as rk
    from product_data
) t where rk = 1

样例运行结果 #

基于你提供的测试数据集，运行后输出如下：

+--------------------+-------+-------+
|       category_code|  brand|  count|
+--------------------+-------+-------+
|electronics.smart...|samsung|1782386|
|electronics.video.tv|samsung| 183988|
|electronics.audio...|  apple| 165277|
|  computers.notebook|   acer| 154599|
|  electronics.clocks|  casio| 141403|
+--------------------+-------+-------+

内容的提问来源于stack exchange，提问作者iblend171