获取数据框行的范围

Question

所以我从 parquet 文件加载了一个数据框。此数据框现在包含未指定数量的列。第一列是Label，后面是features

我想将数据框中的每一行保存为一个 LabeledPoint。

到目前为止我在想：

val labeledPoints: RDD[LabeledPoint] =df.map{row => LabeledPoint(row.getInt(0),Vectors.dense(row.getDouble(1),row.getDouble(2)))}

获取列索引很容易，但是当处理很多列时，这将不适用。我希望能够将从索引 1 开始的整行（因为索引 0 是标签）加载到密集向量中。

有什么想法吗？

Answer 1

这应该可以解决问题

df.map {
  row: Row => 
    val data = for (index <- 1 until row.length) yield row.getDouble(index)
    val vector = new DenseVector(data.toArray)
    new LabeledPoint(row.getInt(0), vector)
}