Skip to main content
 首页 » 编程设计

python之Pyspark之将另一列添加到稀疏向量列

2024年09月07日36lyhabc

我有一个 PySpark 数据框,其中一列 (features) 是一个稀疏向量。例如:

+------------------+-----+ 
|     features     |label| 
+------------------+-----+ 
| (4823,[87],[0.0])|  0.0| 
| (4823,[31],[2.0])|  0.0| 
|(4823,[159],[0.0])|  1.0| 
|  (4823,[1],[7.0])|  0.0| 
|(4823,[15],[27.0])|  0.0| 
+------------------+-----+ 

我想扩展 features 列并向其添加另一个功能,例如:

+-------------------+-----+ 
|     features      |label| 
+-------------------+-----+ 
| (4824,[87],[0.0]) |  0.0| 
| (4824,[31],[2.0]) |  0.0| 
|(4824,[159],[0.0]) |  1.0| 
|  (4824,[1],[7.0]) |  0.0| 
|(4824,[4824],[7.0])|  0.0| 
+-------------------+-----+ 

有没有一种方法可以做到这一点,而无需将 SparseVector 解压缩为密集,然后使用新列将其重新打包为稀疏?

请您参考如下方法:

使用 VectorAssembler 可以最简单地向现有 SparseVector 添加新列ML 库中的转换器。它会自动将列组合成一个向量(DenseVector 或 SparseVector,具体取决于哪个使用的内存最少)。使用 VectorAssembler 将不会在合并过程中将向量转换为 DenseVector(参见 source code)。它可以按如下方式使用:

df = ... 
 
assembler = VectorAssembler( 
    inputCols=["features", "new_col"], 
    outputCol="features") 
 
output = assembler.transform(df) 

要简单地增加 SparseVector 的大小,而不添加任何新值,只需创建一个更大大小的新向量:

def add_empty_col_(v): 
    return SparseVector(v.size + 1, v.indices, v.values) 
 
add_empty_col = udf(add_empty_col_, VectorUDT()) 
df.withColumn("sparse", add_empty_col(col("features"))