Pyspark:对于一列向量计算余弦相似度的最快方法是什么

时间:2018-02-27 22:36:25

标签: python apache-spark dataframe pyspark apache-spark-sql

初学Pyspark问题在这里!我有一个~2M行已经矢量化文本的数据帧(通过w2v; 300维度)。 针对新的单个矢量输入计算每行的余弦距离的最有效方法是什么?

我目前的方法使用了一个udf,需要花费几分钟,对于我想要创建的webapp来说太长了。

创建示例df:

import numpy as np
import pandas as pd
from pyspark.sql.functions import *

column=[]
num_rows = 10000 #change to 2000000 to really slow your computer down!
for x in range(num_rows):
    sample = np.random.uniform(low=-1, high=1, size=(300,)).tolist()
    column.append(sample)
index = range(1000)
df_pd = pd.DataFrame([index, column]).T
#df_pd = pd.concat([df.T[x] for x in df.T], ignore_index=True)
df_pd.head()
df = spark.createDataFrame(df_pd).withColumnRenamed('0', 'Index').withColumnRenamed('1', 'Vectors')
df.show()

创建一个示例输入(我将其创建为spark df以便转换现有管道):

new_input = np.random.uniform(low=-1, high=1, size=(300,)).tolist()
df_pd_new = pd.DataFrame([[new_input]])
df_new = spark.createDataFrame(df_pd_new, ['Input_Vector'])
df_new.show()

计算Vector和new_input之间的余弦距离或相似度:

value = df_new.select('Input_Vector').collect()[0][0]
def cos_sim(vec):
    if (np.linalg.norm(value) * np.linalg.norm(vec)) !=0:
        dot_value = np.dot(value, vec) / (np.linalg.norm(value)*np.linalg.norm(vec))
        return dot_value.tolist()
cos_sim_udf = udf(cos_sim, FloatType())

#df_all_cos = df_all.withColumn('cos_dis', dot_product_udf('w2v')).dropna(subset='cos_dis')
df_cos = df.withColumn('cos_dis', cos_sim_udf('Vectors')).dropna(subset='cos_dis')
df_cos.show()

最后让我们拿出最多5个指数以获得乐趣:

max_values = df_cos.select('index','cos_dis').orderBy('cos_dis', ascending=False).limit(5).collect()
top_indicies = []
for x in max_values:
    top_indicies.append(x[0])
print top_indicies

没有余弦距离的pyspark函数(这将是理想的),所以我不确定如何加快速度。任何想法都非常感激!

0 个答案:

没有答案