如何向PySpark DataFrame添加稀疏向量列及现有代码bug排查

原代码Bug清单 #

• 分区字段错误：原始数据主键列是id，你窗口分区写的user_id属于不存在的列，会直接触发字段不存在报错
• UDF传参错误：UDF不能直接接收整个DataFrame作为输入参数，你调用时传入y1（完整DataFrame对象）属于完全错误的传参方式
• 逻辑冗余性能差：窗口函数会给同分组内的每一行都生成重复的向量结果，后续还需要额外去重，而且UDF内部转Pandas、调用scipy的计算逻辑在大数据量下性能极低
• 类型兼容问题：如果后续向量要用于SparkML链路，scipy生成的数组结构兼容性远不如PySpark原生的SparseVector类型

修复后的实现代码 #

方案1：最小改动修复原逻辑（保留scipy生成逻辑） #

import numpy as np
from scipy import sparse
from pyspark.sql import functions as F
from pyspark.sql.types import ArrayType, FloatType

# 先按id、timestamp分组，把同组的v_row、v_col、v_val聚合为数组
y_grouped = y1.groupBy("id", "timestamp").agg(
    F.collect_list("v_row").alias("rows"),
    F.collect_list("v_col").alias("cols"),
    F.collect_list("v_val").alias("vals")
)

# 重写UDF，直接接收数组参数，不需要转Pandas
def sparse_vector(rows, cols, vals):
    V = sparse.coo_matrix((vals, (rows, cols)), shape = (1, 8)).toarray().tolist()
    return V

sparse_vector_udf = F.udf(sparse_vector, ArrayType(ArrayType(FloatType())))

# 生成最终结果
y2 = y_grouped.withColumn("sparse_vector", sparse_vector_udf("rows", "cols", "vals")) \
              .select("id", "timestamp", "sparse_vector")

方案2：纯PySpark原生实现（无第三方依赖，性能更高） #

不需要依赖scipy和numpy，直接用PySpark内置函数生成你要的嵌套数组结构：

from pyspark.sql import functions as F

y2 = y1.groupBy("id", "timestamp").agg(
    F.map_from_entries(F.collect_list(F.struct("v_col", "v_val"))).alias("col_val_map")
).select(
    "id",
    "timestamp",
    # 要和你预期一致的双层嵌套结构，所以外面多包一层array
    F.array(F.array([F.coalesce(F.col("col_val_map").getItem(i), F.lit(0.0)) for i in range(8)])).alias("sparse_vector")
)

输出验证 #

运行后得到的结果和预期完全一致：

内容的提问来源于stack exchange，提问作者Liyuan Zhang