PySpark中将逗号分隔字符串列转换为稠密向量的方法

问题背景 #

现有两个结构一致的PySpark DataFrame，需执行cross join后计算余弦相似度，其中qry_emb为字符串类型列，存储逗号分隔的嵌入向量数值，DataFrame表结构如下：

df.printSchema()
# root
# |-- query: string (nullable = true)
# |-- qry_emb: string (nullable = true)

DataFrame示例参考：

字符串转Dense Vector实现 #

可选择以下任意一种方式完成转换：

• 内置函数转换（推荐，无UDF序列化开销，性能更好，适配Spark 3.0+版本）

from pyspark.sql import functions as F

# 对第一个DataFrame做转换
df1 = df1.withColumn(
    "qry_emb_vec",
    # 先按逗号拆分字符串为数组，转浮点类型后调用内置函数转为Dense Vector
    F.array_to_vector(F.split(F.col("qry_emb"), ",").cast("array<float>"))
)

# 提前重命名第二个DataFrame的同名列，避免cross join后字段冲突
df2 = df2.withColumnRenamed("query", "doc_query").withColumnRenamed("qry_emb", "doc_emb")
df2 = df2.withColumn(
    "doc_emb_vec",
    F.array_to_vector(F.split(F.col("doc_emb"), ",").cast("array<float>"))
)

• UDF转换（逻辑直观，兼容所有Spark版本）

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

# 定义转换UDF
@F.udf(returnType=VectorUDT())
def str_to_densevec(emb_str):
    if not emb_str:
        return None
    return DenseVector([float(val.strip()) for val in emb_str.split(",")])

# 分别对两个DataFrame执行转换
df1 = df1.withColumn("qry_emb_vec", str_to_densevec("qry_emb"))
df2 = df2.withColumnRenamed("query", "doc_query").withColumnRenamed("qry_emb", "doc_emb")
df2 = df2.withColumn("doc_emb_vec", str_to_densevec("doc_emb"))

Cross Join与余弦相似度计算 #

完成向量转换后即可执行关联和相似度计算：

# 执行笛卡尔关联
cross_res = df1.crossJoin(df2)

# Spark 3.2+可直接用内置余弦相似度函数计算
final_res = cross_res.withColumn(
    "cosine_similarity",
    F.cosine_similarity(F.col("qry_emb_vec"), F.col("doc_emb_vec"))
)

如果使用低于3.2的Spark版本，可通过UDF实现相似度计算：

from pyspark.sql.types import FloatType
import numpy as np

@F.udf(returnType=FloatType())
def cos_sim(v1, v2):
    return float(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)))

final_res = cross_res.withColumn("cosine_similarity", cos_sim("qry_emb_vec", "doc_emb_vec"))

注意：Cross join会生成笛卡尔积，当两个DataFrame数据量较大时会出现数据量爆炸，极易耗尽集群计算资源，大数据量场景建议使用LSH算法做近似近邻匹配替代全量cross join计算。

内容的提问来源于stack exchange，提问作者Sree Aurovindh