ROW_NUMBER（TopN）

Top10的示例如下。

SELECT  main_id, type
FROM (
    SELECT  main_id
           ,type
           ,ROW_NUMBER() OVER(PARTITION BY main_id ORDER BY type DESC) AS rn
    FROM   data_demo2
) A
WHERE   A.rn <= 10;

当发生数据倾斜时，可以通过以下几种方式解决：

方案一：SQL写法的两阶段聚合。

为使Map阶段中Partition各分组数据尽可能均匀，增加随机列，将其作为Partition中一参数。

-- 1. 随机数基于源表行分布
SELECT  main_id, type
FROM (
    SELECT  main_id, type
           ,ROW_NUMBER() OVER(PARTITION BY main_id ORDER BY type DESC) AS rn
    FROM (
        SELECT  main_id, type
        FROM (
            SELECT  main_id
                   ,type
                   ,ROW_NUMBER() OVER(PARTITION BY main_id, src_pt ORDER BY type DESC) AS rn
            FROM (
                SELECT  main_id, type
                       ,CAST(ceil(110 * rand()) % 11 AS INT) AS src_pt
                FROM   data_demo2
            )
        ) B
        WHERE B.rn <= 10
    )
) A
WHERE A.rn <= 10;
-- 2. 随机数直接自定义
SELECT  main_id, type
FROM (
    SELECT  main_id, type
           ,ROW_NUMBER() OVER(PARTITION BY main_id ORDER BY type DESC) AS rn
    FROM (
        SELECT  main_id, type
        FROM (
            SELECT  main_id, type
                   ,ROW_NUMBER() OVER(PARTITION BY main_id, src_pt ORDER BY type DESC) AS rn
            FROM (
                SELECT  main_id, type
                      , CAST(ceil(10 * rand()) AS INT) AS src_pt
                FROM   data_demo2
            )
        ) B
        WHERE B.rn <= 10
    )
) A
WHERE A.rn <= 10;

方案二：UDAF写法的两阶段聚合。

Spark推荐用窗口函数 + approx_count_distinct或collect_set限制，或直接借助groupBy().agg(collect_list()) + UDF取 TopN。

# 用Spark内置算子组合实现（更地道，避免UDAF）
from pyspark.sql import functions as F
from pyspark.sql.window import Window
# Step1: 增加随机列
df1 = df.withColumn("src_pt", F.ceil(F.rand()*10))
# Step2: 在打散后的分区取TopN
w1 = Window.partitionBy("main_id","src_pt").orderBy(F.col("type").desc())
df2 = df1.withColumn("rn", F.row_number().over(w1)).filter("rn <= 10")
# Step3: 合并后再次取TopN
w2 = Window.partitionBy("main_id").orderBy(F.col("type").desc())
result = df2.withColumn("rn2", F.row_number().over(w2)).filter("rn2 <= 10").drop("rn","rn2","src_pt")