Pyspark UDF比较稀疏向量

Question

I am trying to write a pyspark UDF that will compare two Sparse Vectors for me. What I would like to write is:

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')))

However, when I run this code I get the following error message:

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)

I have run into similar problems to this before which relate to the dataframe type being unable to deal with numpy data types. Previously I have been able to solve these issues by coercing the numpy array into a list before returning it, but in this case it seems that I am unable to even pull the data out of the SparseVector, for example even the following does not work:

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

compare_udf = udf(compare, IntegerType())

I have been able to circumvent the issues using an RDD but I am still find this a frustrating limitation with the pyspark UDF. Any advice or help appreciated!

Answer 1

I seem to have solved this problem myself - the issue comes down to the fact that the constituent components of the mllib Sparse Vector are numpy types, which are themselves not supported by the pyspark DataFrame. The following adjusted code works:

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')))

While this works it seems somewhat counter-intuitive to me that the pyspark DataFrame would support a constructed datatype (SparseVector) but not it's constituent parts by themselves (numpy integers) nor provide a more enlightening error message explaining the problem.