将 pyspark dataframe 中的多列转换为一个字典

Question

我创建了一个像这样的 PySpark DataFrame：'''

df = spark.createDataFrame([
    ('v', 3, 'a'),
    ('d', 2, 'b'),
    ('q', 9, 'c')],
    ["c1", "c2", "c3"]
)

df.show()

c1  | c2 | c3
v   |  3 | a
d   |  2 | b
q   |  9 | c

我想创建一个这样的新列：

+--------------------------+
|           c4             |
+--------------------------+
|{"c1":"v","c2":3,"c3":"a"}|
|{"c1":"d","c2":2,"c3":"b"}|
|{"c1":"q","c2":9,"c3":"c"}|
+--------------------------+

我希望c4是MapType类型，而不是StringType另外，我想保持值的类型不变。 （将 3,2 和 9 保留为整数，而不是字符串）

Answer 1

如果要获取 JSON 字符串，请像这样使用struct + to_json ：

import pyspark.sql.functions as F

df1 = df.select(
    F.to_json(
        F.struct(*[F.col(c) for c in df.columns])
    ).alias("c4")
)

df1.show(truncate=False)
#+--------------------------+
#|c4                        |
#+--------------------------+
#|{"c1":"v","c2":3,"c3":"a"}|
#|{"c1":"d","c2":2,"c3":"b"}|
#|{"c1":"q","c2":9,"c3":"c"}|
#+--------------------------+

---

编辑

如果您想要 MapType 列，请使用create_map function：

from itertools import chain

df1 = df.select(
    F.create_map(
        *list(chain(*[[F.col(c), F.lit(c)] for c in df.columns]))
    ).alias("c4")
)

#+---------------------------+
#|c4                         |
#+---------------------------+
#|{v -> c1, 3 -> c2, a -> c3}|
#|{d -> c1, 2 -> c2, b -> c3}|
#|{q -> c1, 9 -> c2, c -> c3}|
#+---------------------------+