python - pyspark groupby 并创建包含其他列字典的列
问题描述
我有这个 pyspark 数据框
df = spark.createDataFrame([("a", "b", "v1", 1234, 56, 78, 9), ("a", "b", "v2", 987, 6, 543, 21), ("c", "d", "v1", 12, 345, 6, 789), ("c", "d", "v2", 9, 876, 5, 4321)], ("k1", "k2", "k3", "ca", "pa", "cb", "pb"))
df.show()
+---+---+---+----+---+---+----+
| k1| k2| k3| ca| pa| cb| pb|
+---+---+---+----+---+---+----+
| a| b| v1|1234| 56| 78| 9|
| a| b| v2| 987| 6|543| 21|
| c| d| v1| 12|345| 6| 789|
| c| d| v2| 9|876| 5|4321|
+---+---+---+----+---+---+----+
基本上我想要做的是通过对前两个键进行分组来转换这个数据框,k1
并将k2
第三个键k3
用作字典的主键,该字典将包含其他列(、、、、)的值ca
,并将pa
包含在新列。这种转换将导致数据框看起来完全像这样:cb
pb
+---+---+--------------------------------------------------------------------------------------------------+
|k1 |k2 |k3 |
+---+---+--------------------------------------------------------------------------------------------------+
|c |d |{"v1": {"pa": 345, "pb": 789, "ca": 12, "cb": 6}, "v2": {"pa": 876, "pb": 4321, "ca": 9, "cb": 5}}|
|a |b |{"v1": {"pa": 56, "pb": 9, "ca": 1234, "cb": 78}, "v2": {"pa": 6, "pb": 21, "ca": 987, "cb": 543}}|
+---+---+--------------------------------------------------------------------------------------------------+
为此,我编写了以下代码,但我认为可以改进此代码(使用 pandas_udf 或其他东西),但我没有设法找到更好的解决方案,我正在寻找任何可能导致更优雅的建议/指导一个有效的解决方案。
def reoganize_col(list_json):
col_data = {}
for d in list_json:
print(d)
for k,v in d.items():
col_data[k] = v
return json.dumps(col_data)
udf_reoganize_col = F.udf(reoganize_col, T.StringType())
df = df.withColumn('x', F.create_map(F.lit('ca'), F.col('ca'),
F.lit('cb'), F.col('cb'),
F.lit('pa'), F.col('pa'),
F.lit('pb'), F.col('pb')))
.groupby(['k1', 'k2']).agg(F.collect_list(F.create_map(F.col('k3'), F.col('x'))).alias('k3'))
df = df.withColumn('k3', udf_reoganize_col(F.col('k3')))
解决方案
您的解决方案几乎就在那里。我建议您使用to_json
而不是 UDF 来提高性能,并使用struct
而不是map
使代码更干净。
(df
.groupBy('k1', 'k2')
.agg(F.collect_list(F.struct('k3', F.struct('pa', 'pb', 'ca', 'cb'))).alias('k3'))
.withColumn('k3', F.to_json(F.map_from_entries('k3')))
.show(10, False)
)
# Output
# +---+---+---------------------------------------------------------------------------------+
# |k1 |k2 |k3 |
# +---+---+---------------------------------------------------------------------------------+
# |c |d |{"v1":{"pa":345,"pb":789,"ca":12,"cb":6},"v2":{"pa":876,"pb":4321,"ca":9,"cb":5}}|
# |a |b |{"v1":{"pa":56,"pb":9,"ca":1234,"cb":78},"v2":{"pa":6,"pb":21,"ca":987,"cb":543}}|
# +---+---+---------------------------------------------------------------------------------+
推荐阅读
- python - 如何使用正则表达式正确解析复杂字符串中的元素
- javascript - Can you set seperate colors for different axes in Chart.js?
- prolog - 如何使用算术运算符在数据库中查找特定值
- javascript - 如何在 n:m 关联中进行单向限制?
- matlab - 如何使用批处理文件打开已打开的 MATLAB 文件?
- ubuntu - 在 Ubuntu 上,如何将 ~/.dotnet/tools 添加到 PATH 变量中?它不在 /etc/environment 或 ~/.bashrc 中,但它仍然在我的 PATH 中
- raku - 哈希中的接受与智能匹配:有什么区别?
- r - knitr opts_knit$set(root.dir = "../") 当 chunk_output_type: 控制台被忽略
- .net - System.Speech.Synthesis 的事件信息中的流编号在哪里?
- java - 如何在 API 28 中设置 android 日期和时间选择器的样式