Scala+Spark3.1实现两表向量点积关联的高效方案咨询

问题背景 #

现有两个Spark数据集：

用户表 #

+-------+---------+
|user_id|    value|
+-------+---------+
|  user1|[1, 2, 3]|
|  user2|[4, 5, 6]|
|  user3|[7, 8, 9]|
+-------+---------+

物品表 #

+---------------+---------------+---------------+
|          item1|          item2|          item3|
+---------------+---------------+---------------+
|[0.5, 0.6, 0.7]|[0.2, 0.3, 0.4]|[0.1, 0.8, 0.9]|
+---------------+---------------+---------------+

需要计算每个用户向量与每个物品向量的点积，生成如下结果：

+-------+-----+-----+-----+
|user_id|item1|item2|item3|
+-------+-----+-----+-----+
|  user1|  3.8|    2|  4.4|
|  user2|  9.2|  4.7|  9.8|
|  user3| 14.6|  7.4| 15.2|
+-------+-----+-----+-----+

最初考虑用cross join后逐行列计算，但该方法低效且内存占用高，需要更优实现方案。

最优实现方案 #

核心思路是避免全量笛卡尔积，利用Spark内置向量运算函数+广播小数据集（物品表仅一行，适合广播）来高效计算。

步骤1：导入必要依赖 #

import org.apache.spark.sql.functions._
import org.apache.spark.ml.linalg.{Vector, Vectors}
import org.apache.spark.sql.types._

步骤2：处理物品表并广播向量映射 #

物品表数据量极小，先将其转换成键值对映射（物品名→向量），再通过广播变量分发到各Executor，避免重复传输：

// 假设物品表已加载为DataFrame itemDF
val itemMap = itemDF.select(collect_list(struct(itemDF.columns.map(col): _*))).first()
  .getAs[Seq[Row]](0).head
  .getValuesMap[Seq[Double]](itemDF.columns)
  .mapValues(v => Vectors.dense(v.toArray))

// 广播物品向量映射
val broadcastItemMap = spark.sparkContext.broadcast(itemMap)

步骤3：定义点积计算UDF #

自定义UDF接收用户向量，遍历广播的物品向量计算点积，返回包含所有物品得分的Row：

val calculateDotProducts = udf((userVector: Seq[Double]) => {
  val userVec = Vectors.dense(userVector.toArray)
  val itemMap = broadcastItemMap.value
  Row.fromSeq(
    itemMap.keys.toList.sorted.map(itemName => userVec.dot(itemMap(itemName)))
  )
})

步骤4：生成结果数据集 #

将用户表的向量列传入UDF，再拆分UDF输出为对应物品列：

// 获取排序后的物品列名
val sortedItemNames = broadcastItemMap.value.keys.toList.sorted

// 假设用户表已加载为DataFrame userDF
val resultDF = userDF.withColumn("dotProducts", calculateDotProducts(col("value")))
  .select(
    col("user_id") +: sortedItemNames.zipWithIndex.map { case (itemName, idx) =>
      col("dotProducts").getItem(idx).alias(itemName)
    }: _*
  )

// 输出结果
resultDF.show()

方案优势 #

• 彻底避免cross join带来的数据膨胀，内存占用仅为用户表+广播的极小物品映射
• 借助Spark MLlib的优化向量运算实现点积，比手动循环计算效率更高
• 扩展性强：新增物品列无需修改核心计算逻辑，仅需保证物品表结构一致

内容的提问来源于stack exchange，提问作者David R