如何在 PySpark ML 中找到向量的 argmax答案

【问题标题】：How to find the argmax of a vector in PySpark ML如何在 PySpark ML 中找到向量的 argmax
【发布时间】：2020-07-23 17:59:54
【问题描述】：

我的模型输出了一个 DenseVector 列，我想找到 argmax。 This page 建议这个函数应该可用，但我不确定语法应该是什么。

是df.select("mycolumn").argmax()吗？

【问题讨论】：

标签： apache-spark pyspark apache-spark-ml

【解决方案1】：

我在 python 中找不到 argmax 操作的文档。但是您可以通过将它们转换为数组来完成它们

对于 pyspark 3.0.0

from pyspark.ml.functions import vector_to_array    
tst_arr = tst_df.withColumn("arr",vector_to_array(F.col('vector_column')))
tst_max=tst_arr.withColumn("max_value",F.array_max("arr"))
tst_max_exp = tst_max.select('*',F.posexplode("arr"))
tst_fin = tst_max_exp.where('col==max_value')

对于 pyspark

from pyspark.sql.functions import udf
@udf
def vect_argmax(row):
    row_arr = row.toArray()
    max_pos = np.argmax(row_arr)
    return(int(max_pos))
tst_fin = tst_df.withColumn("argmax",vect_argmax(F.col('probability')))

【讨论】：

我最终做了一些类似于你编写的第二个脚本的事情——我只是怀疑 Scala/Java 方法必须以某种方式在 Python 中公开，但我不知道具体如何
很高兴它成功了。如果这有帮助，请考虑接受并投票以关闭线程。

【解决方案2】：

你试过了吗

from pyspark.sql.functions import col
df.select(col("mycolumn").argmax())

【讨论】：