我想用Pandas UDF计算两个向量的余弦相似度.我用Spark UDF实现了它，它与以下脚本一起工作得很好.
import numpy as np
from pyspark.sql.functions import udf
from pyspark.sql.types import FloatType
# Create dataframe
df = spark.createDataFrame([("A", [1, 2, 3], [3, 4, 5]), ("B", [5, 6, 7], [7, 8, 9] )], ("name", "vec1", "vec2"))
# Cosime Similarity function
def cosine_similarity(vec1, vec2):
return float(np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2)))
# Spark UDF
cosine_similarity_udf = udf(cosine_similarity, FloatType())
当我用Pandas UDF包装它时，如下所示，它给我一个TypeError，显示为TypeError: only size-1 arrays can be converted to Python scalars
import pandas as pd
from pyspark.sql.functions import pandas_udf, PandasUDFType
@pandas_udf(returnType=FloatType())
def cosine_similarity_udf(vec1: pd.Series, vec2: pd.Series) -> pd.Series:
return pd.Series(cosine_similarity(vec1, vec2))
What should be the correct way to get this desired output using Pandas UDF?
