PySpark DataFrame 中的向量列上的 UDF 问题

Question

我在 PySpark 中的向量列上使用 UDF 时遇到问题，可以在此处说明：

from pyspark import SparkContext
from pyspark.sql import Row
from pyspark.sql.types import DoubleType
from pyspark.sql.functions import udf
from pyspark.mllib.linalg import Vectors

FeatureRow = Row('id', 'features')
data = sc.parallelize([(0, Vectors.dense([9.7, 1.0, -3.2])),
                       (1, Vectors.dense([2.25, -11.1, 123.2])),
                       (2, Vectors.dense([-7.2, 1.0, -3.2]))])
df = data.map(lambda r: FeatureRow(*r)).toDF()

vector_udf = udf(lambda vector: sum(vector), DoubleType())

df.withColumn('feature_sums', vector_udf(df.features)).first()

这会失败并显示以下堆栈跟踪：

Py4JJavaError: An error occurred while calling z:org.apache.spark.api.python.PythonRDD.collectAndServe.
: org.apache.spark.SparkException: Job aborted due to stage failure: Task 5 in stage 31.0 failed 1 times, most recent failure: Lost task 5.0 in stage 31.0 (TID 95, localhost): org.apache.spark.api.python.PythonException: Traceback (most recent call last):
  File "/Users/colin/src/spark/python/lib/pyspark.zip/pyspark/worker.py", line 111, in main
    process()
  File "/Users/colin/src/spark/python/lib/pyspark.zip/pyspark/worker.py", line 106, in process
    serializer.dump_stream(func(split_index, iterator), outfile)
x1  File "/Users/colin/src/spark/python/lib/pyspark.zip/pyspark/serializers.py", line 263, in dump_stream
    vs = list(itertools.islice(iterator, batch))
  File "/Users/colin/src/spark/python/pyspark/sql/functions.py", line 469, in <lambda>
    func = lambda _, it: map(lambda x: f(*x), it)
  File "/Users/colin/pokitdok/spark_mapper/spark_mapper/filters.py", line 143, in <lambda>
TypeError: unsupported operand type(s) for +: 'int' and 'NoneType'

查看传递给 UDF 的内容，似乎有些奇怪。传递的参数应该是 Vector，但它会传递一个 Python 元组，如下所示：

(1, None, None, [9.7, 1.0, -3.2])

难道不能在 Vectors 的 DataFrame 列上使用 UDF 吗？

编辑

所以邮件列表中指出这是known issue。将接受来自@hyim 的答案，因为它确实为密集向量提供了临时解决方法。

Answer 1

在 spark-sql 中，向量被处理为（类型、大小、索引、值）元组。

您可以在带有 pyspark 的向量上使用 udf。只需修改一些代码以使用向量类型的值。

vector_udf = udf(lambda vector: sum(vector[3]), DoubleType())

df.withColumn('feature_sums', vector_udf(df.features)).first()

https://github.com/apache/spark/blob/master/mllib/src/main/scala/org/apache/spark/mllib/linalg/Vectors.scala