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

你试过了吗

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