PySpark实现：将DataFrame列值统计结果转换为Map类型

问题描述 #

给定以下PySpark DataFrame：

df = spark.createDataFrame(
    [('new',), ('new',), ('old',)],
    ["x"]
)

需要输出一个仅含counts列的DataFrame，该列值为列x中各元素与其计数组成的Map（Schema为MapType(StringType(), IntegerType(), True)）。

之前尝试的代码返回了两行结果，不符合需求：

agg = df.groupBy("x").agg(F.count("*").alias("count"))
agg.select(F.create_map(agg["x"], agg["count"]).alias("counts")).show()

期望输出：

+--------------------+
|              counts|
+--------------------+
|{new -> 2, old -> 1}|
+--------------------+

实现方法 #

通过先分组统计计数，再将所有键值对合并为单个Map即可实现：

import pyspark.sql.functions as F

# 分组统计每个元素的出现次数
agg_df = df.groupBy("x").agg(F.count("*").alias("count"))

# 将所有(x, count)键值对合并为单个Map
result_df = agg_df.agg(
    F.map_from_entries(F.collect_list(F.struct("x", "count"))).alias("counts")
)

result_df.show()

步骤说明 #

1. 分组统计：使用groupBy + count得到每个元素的出现次数，生成包含x和count的中间DataFrame。
2. 构建键值对结构体：用struct将每个元素的x值和对应的count包装成键值对格式。
3. 收集键值对数组：通过collect_list将所有键值对结构体收集为一个数组。
4. 转换为Map：使用map_from_entries将键值对数组直接转换为Map类型的列，最终得到单行的目标结果。

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