在Scala中将数据帧转换为Spark mllib矩阵

我有一个名为

df

的Spark数据帧作为输入：

+---------------+---+---+---+---+
|Main_CustomerID| A1| A2| A3| A4|
+---------------+---+---+---+---+
|            101|  1|  0|  2|  1|
|            102|  0|  3|  1|  1|
|            103|  2|  1|  0|  0|
+---------------+---+---+---+---+

我需要将

A1

、

A2

、

A3

、

A4

的值收集到mllib矩阵中，例如

dm: org.apache.spark.mllib.linalg.Matrix =
1.0  0.0  2.0  1.0
0.0  3.0  1.0  1.0
2.0  1.0  0.0  0.0

---

如何在Scala中实现这一点？

您可以按如下方式执行，首先获取矩阵中应包含的所有列：

import org.apache.spark.sql.functions._

val matrixColumns = df.columns.filter(_.startsWith("A")).map(col(_))

然后将数据帧转换为

RDD[Vector]

。因为向量需要包含双倍体，所以这个转换也需要在这里完成

import org.apache.spark.mllib.linalg.Vectors
import org.apache.spark.mllib.linalg.distributed.{IndexedRow, IndexedRowMatrix}

val rdd = df.select(array(matrixColumns:_*).as("arr")).as[Array[Int]].rdd
  .zipWithIndex()
  .map{ case(arr, index) => IndexedRow(index, Vectors.dense(arr.map(_.toDouble)))}

然后将rdd转换为

IndexedRowMatrix

，如果需要，可以将其转换为局部矩阵：

val dm = new IndexedRowMatrix(rdd).toBlockMatrix().toLocalMatrix()

---

对于可以收集给驾驶员的较小矩阵，有一种更简单的替代方法：

val matrixColumns = df.columns.filter(_.startsWith("A")).map(col(_))

val arr = df.select(array(matrixColumns:_*).as("arr")).as[Array[Int]]
  .collect()
  .flatten
  .map(_.toDouble)

val rows = df.count().toInt
val cols = matrixColumns.length

// It's necessary to reverse cols and rows here and then transpose
val dm = Matrices.dense(cols, rows, arr).transpose()

---

这个答案很有帮助。但是为什么呢.as[Array[Int]]？@PRIYAM:很乐意帮助：）as[]部分只是为了在转换为rdd或收集数据帧时更加清楚。另一种方法是这样做，例如：

df.rdd.map（u.getAs[Array[Int]]（0））

。