Question

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

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

进入Spark ml SparseVector。问题是，为了创建一个SparseVector，我还需要一个输入参数 - 它的大小。在我的应用程序中，我将Vector大小保存在单独的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

这样做：

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的成员，我不会将它与它的元素分开。

Answer 2

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

如果vecSize 是整数常量，那么您只需使用lit 内置函数作为

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

Spark - 为udf提供额外参数

2 个答案: