如何將行傳遞到 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')))