如何将一些 pyspark 数据框的列转换为具有列名的 dict 并将它们组合成 json 列？

Question

我有以下格式的数据，我想使用带有两列（“标签”和“数据”）的 pyspark 更改其格式。 'tag'列值是唯一的，'data'列值是从原始列'date、stock、price'获取的json字符串 其中将“stock”和“price”组合为“A”列值，将“date”和“num”组合为“B”列值。

我没有找到或写出好的函数来实现这个效果。

我的 spark 版本是 2.1.0

原始数据帧

date, stock, price， tag, num
1388534400, GOOG, 50, a, 1
1388534400, FB, 60, b, 2
1388534400, MSFT, 55, c, 3
1388620800, GOOG, 52, d, 4

我期望输出：

新数据框

tag|                               data

'a'| "{'A':{'stock':'GOOD', 'price': 50}, B:{'date':1388534400, 'num':1}"

'b'| "{'A':{'stock':'FB', 'price': 60}, B:{'date':1388534400, 'num':2}"

'c'| "{'A':{'stock':'MSFT', 'price': 55}, B:{'date':1388534400, 'num':3}"

'd'| "{'A':{'stock':'GOOG', 'price': 52}, B:{'date':1388620800, 'num':4}"

+--+------------------------------- -------------------+

from pyspark.sql import SparkSession
from pyspark.sql.functions import create_map
spark = SparkSession.builder.appName("example").getOrCreate()
df = spark.createDataFrame([
(1388534400, "GOOG", 50, 'a', 1),
(1388534400, "FB", 60, 'b', 2),
(1388534400, "MSFT", 55, 'c', 3),
(1388620800, "GOOG", 52, 'd', 4)]
).toDF("date", "stock", "price", 'tag', 'num')
df.show()
tag_cols = {'A':['stock', 'price'], 'B':['date', 'num']}

# todo, change the Dataframe columns format

Answer 1

IIUC，只需使用 pyspark.sql.functions.struct 和 pyspark.sql.functions.to_json（两者都应该在 spark 2.1 中可用）

from pyspark.sql import functions as F

# skip df initialization[enter link description here][1]

df_new = df.withColumn('A', F.struct('stock', 'price')) \
           .withColumn('B', F.struct('date', 'num')) \
           .select('tag', F.to_json(F.struct('A', 'B')).alias('data'))

>>> df_new.show(5,0)
+---+-----------------------------------------------------------------+
|tag|data                                                             |
+---+-----------------------------------------------------------------+
|a  |{"A":{"stock":"GOOG","price":50},"B":{"date":1388534400,"num":1}}|
|b  |{"A":{"stock":"FB","price":60},"B":{"date":1388534400,"num":2}}  |
|c  |{"A":{"stock":"MSFT","price":55},"B":{"date":1388534400,"num":3}}|
|d  |{"A":{"stock":"GOOG","price":52},"B":{"date":1388620800,"num":4}}|
+---+-----------------------------------------------------------------+