如何将行传递到 pyspark udf

Question

我有一个包含列（id、id2、vec、vec2）的表，其中 id 是整数，向量是 pyspark SparseVeectors。 我想写一个 udf 来获取 vec 和 vec2 的点积，如下所示：

def dot_product(vec, vec2): #or do the row and I can access them later as row.vec, row.vec2?
    return vec.dot(vec2)
udf = udf(dot_product, FloatType())
dot_product = df.withColumn('dot_product', udf('vec', 'vec2'))

但这会将字符串 'vec' 和 'vec2' 作为变量传递。 我试过 df.vec 和 df.vec 但它通过了整个专栏。 我怎样才能做到这一点作为一个行操作？ 我不能使用 selectExpr，因为 SQL 没有点 function。 如果有 SQL 方法，我也会对此持开放态度

Answer 1

将您的 function 转换为 pyspark udf

import pyspark.sql.functions as F
dot_prod_udf = F.udf(dotproduct,<returnType>)
df.select(*,dot_prod_udf(F.col('vec'),F.col('vec2')))