如何根据spark中的另一列数据更新列名

时间:2018-01-02 08:50:33

标签: python apache-spark dataframe pyspark

我的要求如下

DF

+-------------+----+------+--------+----+
|          des|  id|veh_cd|veh_p_cd|code|
+-------------+----+------+--------+----+
|BodilyInjured|1001|    45|      50| BI_|
|BodilyInjured|1001|    12|      55| BI_|
|    Collision|1001|    18|      21|col_|
|    Collision|1002|    21|      25|col_|
+-------------+----+------+--------+----+

如果“des”是“BodilyInjured”,则应为“veh_cd”添加两个单独的列“BI_veh_cd”,为“veh_p_cd”添加BI_veh_p_cd,其中应包含veh_cd和veh_p_cd的最大值,“碰撞”< / p>

注意:需要为列名添加相应的“代码”到“veh_cd”和“veh_p_cd”,并且应填充相应的数据。列名格式应为code + columnname。例如:对于BodilyInjured代码,BI现在列名应为“BI_veh_cd”和“BI_veh_p_cd”

预期产出

id  BI_veh_cd   BI_veh_p_cd     col_veh_cd  col_veh_p_cd
1001    45        55            18            21
1002    NULL      NULL         21             25

1 个答案:

答案 0 :(得分:3)

您正在寻找枢轴功能:

from pyspark.sql.functions import max

df1 = df.groupBy("id") \
       .pivot("code") \
       .agg(max("veh_cd"), max("veh_p_cd")) \
       .orderBy("id")

要重命名列,我们可以使用vanilla Python regex:

import re

newnames = [re.sub("max|\(|\)", '', i) for i in df1.columns]
df1.toDF(*newnames).show()
+----+----------+------------+-----------+-------------+
|  id|BI__veh_cd|BI__veh_p_cd|col__veh_cd|col__veh_p_cd|
+----+----------+------------+-----------+-------------+
|1001|        45|          55|         18|           21|
|1002|      null|        null|         21|           25|
+----+----------+------------+-----------+-------------+