python - 用于比较稀疏向量的 Pyspark UDF

Question

我正在尝试编写一个 pyspark UDF 来比较两个稀疏向量。我想写的是:

from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, IntegerType, FloatType

def compare(req_values, values):
    return [req for req in req_values.indices if req not in values.indices]

compare_udf = udf(compare, ArrayType(IntegerType()))

display(data.limit(5).select('*', compare_udf('req_values', 'values').alias('missing')))

但是，当我运行此代码时，我收到以下错误消息:

SparkException: Job aborted due to stage failure: Task 0 in stage 129.0 failed 4 times, most recent failure: Lost task 0.3 in stage 129.0 (TID 1256, 10.139.64.15, executor 2): net.razorvine.pickle.PickleException: expected zero arguments for construction of ClassDict (for numpy.dtype)

我之前遇到过与此类似的问题，这些问题与数据帧类型无法处理 numpy 数据类型有关。以前，我已经能够通过在返回之前将 numpy 数组强制转换为列表来解决这些问题，但在这种情况下，我似乎无法从 SparseVector 中提取数据，例如，即使以下方法也不起作用:

def compare(req_values, values):
    return req_values.indices[0]   

compare_udf = udf(compare, IntegerType())

我已经能够使用 RDD 来规避这些问题，但我仍然发现这是 pyspark UDF 的一个令人沮丧的限制。任何建议或帮助表示赞赏!

Answer 1

我自己似乎已经解决了这个问题 - 问题归结为以下事实:mllib Sparse Vector 的组成部分是 numpy 类型，而 pyspark DataFrame 本身不支持这些类型。以下调整后的代码有效:

from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, IntegerType, FloatType

def compare(req_values, values):
    return [int(req) for req in req_values.indices if req not in values.indices]

compare_udf = udf(compare, ArrayType(IntegerType()))

display(data.limit(5).select('*', compare_udf('req_values', 'values').alias('missing')))

虽然这有效，但对我来说似乎有点违反直觉，pyspark DataFrame 将支持构造的数据类型(SparseVector)，但不支持它本身的组成部分(numpy 整数)，也不会提供更具启发性的错误消息来解释问题。