在 PySpark 的 DataFrame 列中存储 DenseVector

Question

我正在尝试将 DenseVector 存储到新列中的 DataFrame 中。

我尝试了以下代码，但得到了一个 AttributeError 说 'numpy.ndarray' object has no attribute '_get_object_id'。

from pyspark.sql import functions
from pyspark.mllib.linalg import Vectors

df = spark.createDataFrame([{'name': 'Alice', 'age': 1},
                            {'name': 'Bob', 'age': 2}])

vec = Vectors.dense([1.0, 3.0, 2.9])

df.withColumn('vector', functions.lit(vec))

我希望每行存储一个向量以用于计算目的。任何帮助表示赞赏。

[Python 3.7.3，Spark 版本 2.4.3，通过 Jupyter All-Spark-Notebook]

编辑

我尝试按照 Florian 的建议遵循here 的答案，但我无法调整 udf 以接受自定义的预构建向量。

conv = functions.udf(lambda x: DenseVector(x), VectorUDT())
# Same with
# conv = functions.udf(lambda x: x, VectorUDT())

df.withColumn('vector', conv(vec)).show()

我收到此错误：

TypeError: Invalid argument, not a string or column: [1.0,3.0,2.9] of type <class 'pyspark.mllib.linalg.DenseVector'>. For column literals, use 'lit', 'array', 'struct' or 'create_map' function.

Answer 1

您可以将 udf 的创建包装在一个函数中，因此它会返回带有您的向量的 udf。下面给出一个例子，希望对你有帮助！

import pyspark.sql.functions as F
from pyspark.ml.linalg import VectorUDT, DenseVector

df = spark.createDataFrame([{'name': 'Alice', 'age': 1},
                            {'name': 'Bob', 'age': 2}])

def vector_column(x): 
    return F.udf(lambda: x, VectorUDT())()

vec = DenseVector([1.0, 3.0, 2.9])
df.withColumn("vector", vector_column(vec)).show()

输出：

+---+-----+-------------+
|age| name|       vector|
+---+-----+-------------+
|  1|Alice|[1.0,3.0,2.9]|
|  2|  Bob|[1.0,3.0,2.9]|
+---+-----+-------------+