Spark SQL 2.4.0中基于hex_map聚合distinct user_id求和的实现问题

问题描述 #

现有一张包含date_id、user_id、hex_id、hex_map字段的表，需按hex_id聚合：将每个hex_id对应的hex_map字段替换为该hex_map中所有hex_id组的distinct user_id计数之和（例如hex_id A的hex_map[A,B,C]对应1+2+3=6）。

原始表数据 #

+----------+-------+------+---------+
|   date_id|user_id|hex_id|  hex_map|
+----------+-------+------+---------+
|2016-11-01|    100|     A|[A, B, C]|
|2016-11-01|    300|     B|      [B]|
|2016-11-01|    400|     B|      [B]|
|2016-11-01|    100|     C|   [B, C]|
|2016-11-01|    200|     C|   [B, C]|
|2016-11-01|    300|     C|   [B, C]|
+----------+-------+------+---------+

聚合中间结果（每个hex_id的distinct user_id计数） #

+------+---------+---+
|hex_id|  hex_map|cnt|
+------+---------+---+
|     A|[A, B, C]|  1|
|     B|      [B]|  2|
|     C|   [B, C]|  3|
+------+---------+---+

期望最终结果 #

+------+---------+---+
|hex_id|  hex_map|cnt|
+------+---------+---+
|     A|       6 |  1|
|     B|       2 |  2|
|     C|       5 |  3|
+------+---------+---+

原SQL问题分析 #

你之前尝试的SQL存在以下问题：

• sss子句缺少AS关键字，语法错误；
• 关联后求和的字段错误，应该使用关联表的cnt而非原表的cnt；
• 最终输出未保留原始cnt字段，也未正确映射替换后的hex_map值。

正确实现SQL #

WITH hex_user_cnt AS (
    -- 计算每个hex_id对应的distinct user_id数量
    SELECT hex_id, hex_map, COUNT(DISTINCT user_id) AS cnt
    FROM tab
    GROUP BY hex_id, hex_map
),
explode_hex_map AS (
    -- 将hex_map数组拆分为单个hex值，保留原始信息
    SELECT 
        h.hex_id AS original_hex,
        h.hex_map AS original_map,
        h.cnt AS original_cnt,
        explode(h.hex_map) AS mapped_hex
    FROM hex_user_cnt h
),
sum_mapped_cnt AS (
    -- 对每个原始hex_id，求和其hex_map中所有hex对应的cnt
    SELECT 
        original_hex,
        original_map,
        original_cnt,
        SUM(huc.cnt) AS mapped_sum
    FROM explode_hex_map ehm
    JOIN hex_user_cnt huc ON ehm.mapped_hex = huc.hex_id
    GROUP BY original_hex, original_map, original_cnt
)
-- 输出最终结果，替换hex_map为求和值
SELECT 
    original_hex AS hex_id,
    mapped_sum AS hex_map,
    original_cnt AS cnt
FROM sum_mapped_cnt
ORDER BY hex_id;

逻辑说明 #

1. hex_user_cnt：按hex_id和hex_map分组，计算每个组的distinct user_id数量，得到中间结果；
2. explode_hex_map：使用explode函数将每个hex_map数组拆分为多行，同时保留原始的hex_id、hex_map和cnt；
3. sum_mapped_cnt：将拆分后的mapped_hex与hex_user_cnt关联，获取每个mapped_hex对应的cnt，再按原始hex_id分组求和，得到hex_map对应的总和；
4. 最后重命名字段，得到符合要求的最终结果。

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