scala - 如何将 Spark 数据帧列从 Array[Int] 转换为 linalg.Vector？

Question

我有一个数据框 df，如下所示:

+--------+--------------------+
| user_id|        is_following|
+--------+--------------------+
|       1|[2, 3, 4, 5, 6, 7]  |
|       2|[20, 30, 40, 50]    |
+--------+--------------------+

我可以确认它具有架构:

root
 |-- user_id: integer (nullable = true)
 |-- is_following: array (nullable = true)
 |    |-- element: integer (containsNull = true)

我想使用 Spark 的 ML 例程(例如 LDA)对此进行一些机器学习，需要我将 is_following 列转换为 linalg.Vector (不是Scala 向量)。当我尝试通过

执行此操作时

import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.ml.linalg.Vectors

val assembler = new VectorAssembler().setInputCols(Array("is_following")).setOutputCol("features")
val output = assembler.transform(df)

然后我收到以下错误:

java.lang.IllegalArgumentException: Data type ArrayType(IntegerType,true) is not supported.

如果我的解释正确，我就会明白我需要将这里的类型从整数转换为其他类型。 (双数？字符串？)

我的问题是，将此数组转换为能够为 ML 管道正确矢量化的数组的最佳方法是什么？

编辑:如果有帮助，我不必以这种方式构建数据帧。我可以改为:

+--------+------------+
| user_id|is_following|
+--------+------------+
|       1|           2|
|       1|           3|
|       1|           4|
|       1|           5|
|       1|           6|
|       1|           7|
|       2|          20|
|     ...|         ...|
+--------+------------+

Answer 1

将数组转换为 linalg.Vector 并同时将整数转换为 double 的一个简单解决方案是使用 UDF。

使用您的数据框:

val spark = SparkSession.builder.getOrCreate()
import spark.implicits._

val df = spark.createDataFrame(Seq((1, Array(2,3,4,5,6,7)), (2, Array(20,30,40,50))))
  .toDF("user_id", "is_following")

val convertToVector = udf((array: Seq[Int]) => {
  Vectors.dense(array.map(_.toDouble).toArray)
})

val df2 = df.withColumn("is_following", convertToVector($"is_following"))

spark.implicits._ 在此导入以允许使用 $、col() 或 ' 可以代替使用。

打印df2数据框将给出想要的结果:

+-------+-------------------------+
|user_id|is_following             |
+-------+-------------------------+
|1      |[2.0,3.0,4.0,5.0,6.0,7.0]|
|2      |[20.0,30.0,40.0,50.0]    |
+-------+-------------------------+

架构:

root
 |-- user_id: integer (nullable = false)
 |-- is_following: vector (nullable = true)