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 向量的第一个元素？

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

How to access the values of denseVector in PySpark

How to access element of a VectorUDT column in a Spark DataFrame?

Answer 1

更新：

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

因此，避免这种愚蠢错误的一种处理方法是使用 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 数组的元素一样访问元素。

有几种方法可以访问数据框中数组的各个元素。一种是在您的 select 语句中显式调用列 cv_predictions_prod['probability']。通过显式调用该列，您可以对该列执行操作，例如选择数组中的第一个元素。例如：

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

应该可以解决问题。