应用模型后，Pyspark提取转换数据框的概率

Question

在应用了随机图classclassifier进行二进制分类并在数据集上进行预测后，我获得了转换 dataframe df df ，带有标签，预测和概率列。
目标：
我想创建一个新的列" prob_flag"，这是预测标签'1'的概率。它是包含概率的数组的第二个元素(本身就是第一个数组的第三个元素(。

我研究了类似的主题，但在这些主题中没有遇到错误。

df.show()
label   prediction                 probability
  0           0           [1,2,[],[0.7558548984793847,0.2441451015206153]]
  0           0           [1,2,[],[0.5190322149055472,0.4809677850944528]]
  0           1           [1,2,[],[0.4884140358521083,0.5115859641478916]]
  0           1           [1,2,[],[0.4884140358521083,0.5115859641478916]]
  1           1           [1,2,[],[0.40305518381637956,0.5969448161836204]]
  1           1           [1,2,[],[0.40570407426458577,0.5942959257354141]]
# The probability column is VectorUDT and looks like an array of dim 4 that contains probabilities of predicted variables I want to retrieve  
df.schema
StructType(List(StructField(label,DoubleType,true),StructField(prediction,DoubleType,false),StructField(probability,VectorUDT,true)))
# I tried this:
import pyspark.sql.functions as f
df.withColumn("prob_flag", f.array([f.col("probability")[3][1])).show()
"Can't extract value from probability#6225: need struct type but got struct<type:tinyint,size:int,indices:array<int>,values:array<double>>;"

我想创建一个新的列" prob_flag"，这是预测标签'1'的概率。它是包含概率的数组的第二个数字，例如0.24、0.48、0.51、0.51、0.59、0.59。

Answer 1

不幸的是，您无法提取vectorudt的字段，就好像它是arraytype一样。

您必须使用UDF：

from pyspark.sql.types import DoubleType
from pyspark.sql.functions import udf, col
def extract_prob(v):
    try:
        return float(v[1])  # Your VectorUDT is of length 2
    except ValueError:
        return None
extract_prob_udf = udf(extract_prob, DoubleType())
df2 = df.withColumn("prob_flag", extract_prob_udf(col("probability")))