如何在 PySpark ML 中找到向量的 argmax [英] How to find the argmax of a vector in PySpark ML

问题描述

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

My model has output a DenseVector column, and I'd like to find the argmax. This page suggests this function should be available, but I'm not sure what the syntax should be.

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

推荐答案

我在 python 中找不到 argmax 操作的文档.但是你可以通过将它们转换为数组来实现

I could not find the documents for argmax operation in python. but you can do them by converting them to arrays

对于 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 <3.0.0

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')))

这篇关于如何在 PySpark ML 中找到向量的 argmax的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！