python - 将值添加到 PySpark 中的 DenseVector 中

Question

我有一个数据框，我已经将其处理为:

+---------+-------+
| inputs  | temp  | 
+---------+-------+
| [1,0,0] | 12    |
+---------+-------+
| [0,1,0] | 10    |
+---------+-------+
...

inputs 是 DenseVectors 的列。 temp 是一列值。我想在 DenseVector 中附加这些值并创建一列，但我不知道如何开始。对于此所需输出的任何提示:

+---------------+
| inputsMerged  | 
+---------------+
| [1,0,0,12]    | 
+---------------+
| [0,1,0,10]    |
+---------------+
...

编辑:我正在尝试使用 VectorAssembler 方法，但我生成的数组不符合预期。

Answer 1

你可以这样做:

df.show()
+-------------+----+
|       inputs|temp|
+-------------+----+
|[1.0,0.0,0.0]|  12|
|[0.0,1.0,0.0]|  10|
+-------------+----+

df.printSchema()
root
 |-- inputs: vector (nullable = true)
 |-- temp: long (nullable = true)

导入:

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

创建 udf 以合并向量和元素:

concat = F.udf(lambda v, e: Vectors.dense(list(v) + [e]), VectorUDT())

将 udf 应用于输入和临时列:

merged_df = df.select(concat(df.inputs, df.temp).alias('inputsMerged'))

merged_df.show()
+------------------+
|      inputsMerged|
+------------------+
|[1.0,0.0,0.0,12.0]|
|[0.0,1.0,0.0,10.0]|
+------------------+

merged_df.printSchema()
root
 |-- inputsMerged: vector (nullable = true)