如何将稀疏向量作为列添加到Pyspark Dataframe

在PySpark中，稀疏向量通常用于表示具有大量零值的向量，以节省存储空间和提高计算效率。要将稀疏向量作为列添加到DataFrame，你可以使用pyspark.ml.linalg.SparseVector类。以下是将稀疏向量作为列添加到PySpark DataFrame的步骤：

基础概念

稀疏向量：一种数据结构，用于表示大部分元素为零的向量。它只存储非零元素及其索引，从而节省空间。
PySpark DataFrame：一个分布式数据集，类似于传统数据库中的表，但分布在集群的多个节点上。

类型与应用场景

类型：常见的稀疏向量格式包括COO（Coordinate List）、CSR（Compressed Sparse Row）和CSC（Compressed Sparse Column）。
应用场景：自然语言处理（如TF-IDF向量）、推荐系统、图像处理等领域。

示例代码

以下是一个示例，展示如何将稀疏向量作为列添加到PySpark DataFrame：

from pyspark.sql import SparkSession
from pyspark.ml.linalg import SparseVector

# 初始化SparkSession
spark = SparkSession.builder.appName("SparseVectorExample").getOrCreate()

# 创建一个示例DataFrame
data = [(1, [0, 0, 2.0, 0, 3.0]), (2, [1.0, 0, 0, 0, 0]), (3, [0, 0, 0, 4.0, 0])]
columns = ["id", "values"]
df = spark.createDataFrame(data, columns)

# 将列表转换为稀疏向量
def to_sparse_vector(values):
    size = len(values)
    indices = [i for i, v in enumerate(values) if v != 0]
    values = [v for v in values if v != 0]
    return SparseVector(size, indices, values)

# 使用UDF将列表列转换为稀疏向量列
from pyspark.sql.functions import udf
from pyspark.sql.types import VectorType

to_sparse_udf = udf(to_sparse_vector, VectorType())
df_with_sparse_vector = df.withColumn("sparse_vector", to_sparse_udf(df["values"]))

# 显示结果
df_with_sparse_vector.show(truncate=False)