找到一个密集向量的最大值将其设置为1，其余所有设置为0 [pyspark]

Question

我有一个包含密集向量的DataFrame列。我想一个这些矢量的最大值设置为1且取代所有其余为0。 例如：

+------+-------------+ 
| Col1|  vector| 
+------+-------------+ 
|Modali|[1.0,2.0,3.0]| 
|assert|[4.0,9.0,3.0]| 
+------+-------------+ 

我想将其变为：

+------+-------------+ 
| Col1|  vector| 
+------+-------------+ 
|Modali|[0.0,0.0,1.0]| 
|assert|[0.0,1.0,0.0]| 
+------+-------------+ 

注意我正在使用python开发Spark 2.1.0。 非常感谢你提前。

Answer 1

您可以创建一个udf，它接受一个向量并将其二进制化;可以通过简单地检查向量中的值是否等于使用列表理解的最大值来构造二进制器：

from pyspark.ml.linalg import Vectors, VectorUDT 
import pyspark.sql.functions as F 

def max_binarizer(vector): 
    max_val = max(vector)        # maximum value in the vector 
    return Vectors.dense([1 if x == max_val else 0 for x in vector]) # binarize it 

# create a udf for the binarizer 
max_bin_udf = F.udf(max_binarizer, VectorUDT()) 

df.withColumn("vector", max_bin_udf(df["vector"])).show() 
+------+-------------+ 
| Col1|  vector| 
+------+-------------+ 
|Modali|[0.0,0.0,1.0]| 
|assert|[0.0,1.0,0.0]| 
+------+-------------+