在Spark DataFrame中访问向量元素（逻辑回归概率向量）

Question

我在PySpark（ML包）中训练了一个LogisticRegression模型，并且预测结果是一个PySpark DataFrame（cv_predictions）（参见[1]）。probability列（参见[2]）是一个vector类型（参见[3]）。

[1]
type(cv_predictions_prod)
pyspark.sql.dataframe.DataFrame

[2]
cv_predictions_prod.select('probability').show(10, False)
+----------------------------------------+
|probability                             |
+----------------------------------------+
|[0.31559134817066054,0.6844086518293395]|
|[0.8937864350711228,0.10621356492887715]|
|[0.8615878905395029,0.1384121094604972] |
|[0.9594427633777901,0.04055723662220989]|
|[0.5391547673698157,0.46084523263018434]|
|[0.2820729747752462,0.7179270252247538] |
|[0.7730465873083118,0.22695341269168817]|
|[0.6346585276598942,0.3653414723401058] |
|[0.6346585276598942,0.3653414723401058] |
|[0.637279255218404,0.362720744781596]   |
+----------------------------------------+
only showing top 10 rows

[3]
cv_predictions_prod.printSchema()
root
 ...
 |-- rawPrediction: vector (nullable = true)
 |-- probability: vector (nullable = true)
 |-- prediction: double (nullable = true)

我该如何解析PySpark DataFrame的vector，以便创建一个新列，仅提取每个probability向量的第一个元素？

这个问题类似于下面的问题，但是链接中的解决方案对我来说不起作用/不清楚：

如何在PySpark中访问denseVector的值？

如何在Spark DataFrame中访问VectorUDT列的元素？

- user2205916

Answer 1

更新：

似乎在Spark中存在一个bug，阻止你在select语句期间访问密集向量中的单个元素。通常情况下，你应该能够像访问numpy数组一样访问它们，但是尝试运行之前发布的代码时，你可能会收到错误信息pyspark.sql.utils.AnalysisException: "Can't extract value from probability#12;"

因此，避免这个愚蠢的bug的一种方法是使用UDF。与其他问题类似，可以按以下方式定义UDF：

from pyspark.sql.functions import udf
from pyspark.sql.types import FloatType

firstelement=udf(lambda v:float(v[0]),FloatType())
cv_predictions_prod.select(firstelement('probability')).show()

在幕后，这仍然像numpy数组一样访问DenseVector的元素，但它不会像之前那样抛出相同的错误。

由于这个回答得到了很多赞，我觉得我应该划掉这个回答中不正确的部分。

原始回答：密集向量只是numpy数组的包装器。因此，您可以以与访问numpy数组元素相同的方式访问元素。 有几种方法可以访问数据帧中数组的单个元素。其中一种方法是在选择语句中显式调用列cv_predictions_prod ['probability']。通过显式调用列，您可以对该列执行操作，例如选择数组中的第一个元素。例如：

cv_predictions_prod.select(cv_predictions_prod['probability'][0]).show()

应该解决这个问题。

- DavidWayne

---

不行，这样做不起作用。VectorUDT没有表示为ArrayType。 - zero323

从pyspark文档中：由值数组表示的密集向量。我们使用numpy数组进行存储，算术运算将委托给底层的numpy数组。http://spark.apache.org/docs/latest/api/python/pyspark.ml.html#pyspark.ml.linalg.DenseVector。当尝试运行示例代码时，您会收到什么错误？ - DavidWayne

这不是一个 bug。Spark 的 DataFrame 与 Python 对象相差甚远，它不使用 NumPy 在后台运行，除非您明确将其转换为 Python RDD（由 udf 使用的批处理 Python eval），而 VectorUDT 不是本地 SQL 类型，因此它不提供与 ArrayType 等相同的功能。 - zero323

更新的答案可以运行，但旧的答案可以吗？ - dksahuji

也使用了 withColumn。 - Shuai Liu