Spark - 为 udf 提供额外的参数

Question

我正在尝试创建一个 spark scala udf 来转换以下形状的 MongoDB 对象：

Object:
    "1": 50.3
    "8": 2.4
    "117": 1.0

进入 Spark ml SparseVector。 问题是，为了创建一个 SparseVector，我还需要一个输入参数——它的大小。 在我的应用程序中，我将向量大小保存在单独的 MongoDB 集合中。 所以，我定义了以下 UDF 函数：

val mapToSparseVectorUdf = udf {
  (myMap: Map[String, Double], size: Int) => {
    val vb: VectorBuilder[Double] = new VectorBuilder(length = -1)
    vb.use(myMap.keys.map(key => key.toInt).toArray, myMap.values.toArray, size)
    vb.toSparseVector
  }
}

我试图这样称呼它：

df.withColumn("VecColumn", mapToSparseVectorUdf(col("MapColumn"), vecSize)).drop("MapColumn")

但是，我的 IDE 对那个 udf 调用说“不适用”。 有没有办法制作这种可以带额外参数的UDF？

Answer 1

Udf 函数需要将列作为参数传递，而传递的 columns 将通过序列化解析为原始数据类型和反序列化。这就是为什么 udf 函数很昂贵

如果vecSize 是一个整数常量，那么你可以简单地使用lit 内置函数作为

df.withColumn("VecColumn", mapToSparseVectorUdf(col("MapColumn"), lit(vecSize))).drop("MapColumn")

Answer 2

这样就可以了：

def mapToSparseVectorUdf(vectorSize: Int) = udf[Vector, Map[String, Double]](
  (myMap: Map[String, Double]) => {
    val elements = myMap.toSeq.map {case (index, value) => (index.toInt, value)}
    Vectors.sparse(vectorSize, elements)
  }
)

用法：

val data = spark.createDataFrame(Seq(
    ("1", Map("1" -> 50.3, "8" -> 2.4)), 
    ("2", Map("2" -> 23.5, "3" -> 41.2))
)).toDF("id", "MapColumn")

data.withColumn("VecColumn", mapToSparseVectorUdf(10)($"MapColumn")).show(false)

注意：

考虑修复您的 MongoDB 架构！ ;) 大小是 SparseVector 的成员，我不会将它与它的元素分开。