如何根據 PySpark 數據幀的另一列中的值修改列？ F.when 邊緣情況

Question

我想遍歷 pyspark 數據框中的每一行，並根據另一列的內容更改列的值。 我將其更改為的值也基於要更改的列的當前值。

具體來說，我有一列包含 DenseVectors，另一列包含我需要的向量的索引。

或者，我也可以將 DenseVector 替換為 DenseVector 中兩個值中較大的一個。

我主要是嘗試將 F.when() 與 withColumn 結合使用，但是我遇到了 F.when() 的第二個元素的麻煩，因為我想存儲向量的正確索引，但不能直接索引一列。

   a                        b  
1  DenseVector([0.1, 0.9])  1.0
2  DenseVector([0.6, 0.4])  0.0
.
.
.

df = df.withColumn('a', F.when(df.b == 0.0, df.a[0])
                  .otherwise(df.a[1])

Answer 1

我能夠通過遵循在這個問題的答案中找到的信息得出一個解決方案。

Spark 中似乎存在一個錯誤，它不允許您對包含在 Spark 數據幀中的 DenseVectors 進行索引。 這可以通過創建一個用戶定義的函數來解決，該函數像一個 numpy 數組一樣訪問元素。

from pyspark.sql import functions as F
from pyspark.sql import types as T

firstelement=F.udf(lambda v:float(v[0]),T.FloatType())
secondelement=F.udf(lambda v:float(v[1]),T.FloatType())


df = df.withColumn('a', F.when(df['b'] == 0.0, 
     firstelement('a')).otherwise(secondelement('a'))