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

Question

我有一个 pyspark dataframe 如下，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 时，它一直显示错误“作业中止由于......”所以我想问一下是否有任何其他方法也可以得到同样的结果。 谢谢！

Answer 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|
+---+---+---+---+------+