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