Pyspark：用同名的另一列替换行值

Posted 2023-04-15

【中文标题】Pyspark：用同名的另一列替换行值

【英文标题】：Pyspark: replace row value by another column with the same name

【发布时间】：2021-04-07 13:49:26

【问题描述】：

我有一个 pyspark 数据框，如下所示，df

| D1 | D2 | D3 |Out|
| 2  | 4  | 5  |D2 |
| 5  | 8  | 4  |D3 |
| 3  | 7  | 8  |D1 |

并且我想将“out”列的行值替换为同一行中与“out”列的行值具有相同列名的行值。

| D1 | D2 | D3 |Out|Result|
| 2  | 4  | 5  |D2 |4     |
| 5  | 8  | 4  |D3 |4     |
| 3  | 7  | 8  |D1 |3     |

df_lag=df.rdd.map(lambda row: row + (row[row.Out],)).toDF(df.columns + ["Result"])

我已经尝试了上面的代码，它可以获得结果，但是当我尝试保存到 csv 时，它一直显示错误“Job aborted due to......”所以我想问一下是否有任何其他方法也可以获得相同的结果。谢谢！

【参考方案1】：

您可以使用链式when 语句，使用reduce 从列名动态生成：

from functools import reduce
import pyspark.sql.functions as F

df2 = df.withColumn(
    'Result', 
    reduce(
        lambda x, y: x.when(F.col('Out') == y, F.col(y)), 
        df.columns[:-1], 
        F
    )
)

df2.show()
+---+---+---+---+------+
| D1| D2| D3|Out|Result|
+---+---+---+---+------+
|  2|  4|  5| D2|     4|
|  5|  8|  4| D3|     4|
|  3|  7|  8| D1|     3|
+---+---+---+---+------+