如何将 PySpark dataframe 列转换为基于 groupBy 列的字典列表

Question

我正在将 dataframe 列转换为字典列表。

输入 dataframe 有 3 列：

ID  accounts pdct_code
1    100      IN
1    200      CC
2    300      DD
2    400      ZZ
3    500      AA

我需要读取此输入 dataframe 并将其转换为 3 output 行。 output 应如下所示：

ID arrayDict
1  [{“accounts”: 100, “pdct_cd”: ’IN’}, {”accounts”: 200, “pdct_cd”: ’CC’}]

同样，对于 ID“2”，应该有 1 行包含 2 个带有键值对的字典。

我试过这个：

Df1 = df.groupBy("ID").agg(collect_list(struct(col("accounts"), ("pdct_cd"))).alias("array_dict"))

但是 output 不是我想要的，它应该是一个字典列表。

Answer 1

您描述的内容（字典列表）在 Spark 中不存在。 我们有 arrays 而不是列表，我们有结构或映射而不是字典。 由于您没有使用这些术语，因此这将是对我认为您需要的内容的松散解释。

以下将创建 arrays 个字符串。 这些字符串将具有您可能想要的结构。

df.groupBy("ID").agg(F.collect_list(F.to_json(F.struct("accounts", "pdct_code")))

struct()将您的列放入结构数据类型中。
to_json()从提供的结构中创建一个 JSON 字符串。
collect_list()是一个聚合 function ，它将组中的所有字符串移动到一个数组中。

完整示例：

from pyspark.sql import functions as F
df = spark.createDataFrame(
    [(1, 100, "IN"),
     (1, 200, "CC"),
     (2, 300, "DD"),
     (2, 400, "ZZ"),
     (3, 500, "AA")],
    ["ID", "accounts", "pdct_code"])

df = df.groupBy("ID").agg(F.collect_list(F.to_json(F.struct("accounts", "pdct_code"))).alias("array_dict"))

df.show(truncate=0)
# +---+----------------------------------------------------------------------+
# |ID |array_dict                                                            |
# +---+----------------------------------------------------------------------+
# |1  |[{"accounts":100,"pdct_code":"IN"}, {"accounts":200,"pdct_code":"CC"}]|
# |3  |[{"accounts":500,"pdct_code":"AA"}]                                   |
# |2  |[{"accounts":300,"pdct_code":"DD"}, {"accounts":400,"pdct_code":"ZZ"}]|
# +---+----------------------------------------------------------------------+