如何计算PySpark Map类型列各元素计数占总计数的百分比

问题场景 #

现有结构如下的PySpark DataFrame，其中Item列为Map<Str,int>类型：

Date              Item (Map<Str,int>)                           Total Items   

2021-02-01    Item_A -> 3, Item_B -> 10, Item_C -> 2                 15
2021-02-02    Item_A -> 1, Item_B -> 5, Item_C ->  7                 13
2021-02-03    Item_A -> 8, Item_B -> 3, Item_C ->  1                 12

需要新增一列，对Map内每个物品按(单个物品计数/总物品数)*100计算计数占比，结果列同样为Map类型，期望输出如下：

Date              Item (Map<Str,int>)                    Total Items     Item count % 
                                                                        (item/total items)*100

2021-02-01    Item_A -> 3, Item_B -> 10, Item_C -> 5           15     Item_A -> 20%, Item_B -> 66%, Item_c -> 33%
2021-02-02    Item_A -> 1, Item_B -> 5, Item_C ->  7           13    Item_A -> 7%, Item_B -> 38%, Item_C -> 53%
2021-02-03    Item_A -> 8, Item_B -> 3, Item_C ->  1           12    Item_A -> 66%, Item_B -> 25%, Item_C -> 8.3%

当前使用的代码无法得到正确结果：

df = df.withColumn('Item_count_percentage', F.expr('aggregate(map_values(Item), 0 , (acc, x) -> (acc / int(x)/100)'))
df.show(truncate=False)

原有代码问题 #

• aggregate函数的作用是对数组元素做迭代聚合，最终只会返回单个聚合值，无法生成键值对结构的Map结果，完全不符合输出类型要求
• 计算逻辑错误，公式写反，没有关联Total Items列参与运算
• 表达式括号未闭合，存在语法错误

正确实现代码 #

通过map_entries拆解Map为键值对结构体数组，遍历计算每个键的占比后再重组为Map即可：

from pyspark.sql import functions as F

df = df.withColumn(
    "Item_count_percentage",
    F.expr("""
        map_from_entries(
            transform(
                map_entries(Item),
                entry -> struct(
                    entry.key,
                    concat(round(entry.value / `Total Items` * 100, 1), '%')
                )
            )
        )
    """)
)

df.show(truncate=False)

逻辑说明 #

• map_entries(Item)：将Map列转换为结构体数组，每个结构体包含key（物品名称）、value（对应物品计数）两个属性
• transform：遍历结构体数组，对每个键值对计算占比：用物品计数除以总物品数乘100，按需要保留小数位后拼接%后缀
• map_from_entries：将处理完成的结构体数组重新转换为Map类型，满足输出列的类型要求
• 注意列名Total Items包含空格，在SQL表达式中需要用反引号包裹，否则会触发语法解析错误

如果需要数值类型的占比结果而非带%的字符串，去掉concat函数，直接保留round计算的数值即可。

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