WrappedArray 的 Spark Dataframe 到 Dataframe[Vector]

Question

我有一个 spark Dataframe df 具有以下架构：

root
 |-- features: array (nullable = true)
 |    |-- element: double (containsNull = false)

我想创建一个新的数据框，其中每一行都是Doubles 的向量，并希望获得以下架构：

root
     |-- features: vector (nullable = true)

到目前为止，我有以下代码（受这篇文章的影响：Converting Spark Dataframe(with WrappedArray) to RDD[labelPoint] in scala），但我担心它有问题，因为即使计算合理数量的行也需要很长时间。 此外，如果行太多，应用程序将因堆空间异常而崩溃。

val clustSet = df.rdd.map(r => {
          val arr = r.getAs[mutable.WrappedArray[Double]]("features")
          val features: Vector = Vectors.dense(arr.toArray)
          features
          }).map(Tuple1(_)).toDF()

我怀疑指令 arr.toArray 在这种情况下不是一个好的 Spark 实践。任何澄清都会非常有帮助。

谢谢！

Answer 1

这是因为.rdd 必须从内部内存格式中反序列化对象，这非常耗时。

可以使用.toArray - 您在行级别上操作，而不是将所有内容收集到驱动程序节点。

您可以使用 UDF 轻松做到这一点：

import org.apache.spark.ml.linalg._
val convertUDF = udf((array : Seq[Double]) => {
  Vectors.dense(array.toArray)
})
val withVector = dataset
  .withColumn("features", convertUDF('features))

代码来自这个答案：Convert ArrayType(FloatType,false) to VectorUTD

但是问题的作者没有询问差异