转换Spark DataFrame以进行ML处理

Question

我编写了以下代码来将数据提供给Spark 2.3中的机器学习算法。下面的代码运行正常。我需要增强此代码，以便能够转换不仅仅3列，而是通过csv文件上传的任意数量的列。例如，如果我加载了5列，如何将它们自动放入下面的Vector.dense命令中，或者以其他方式生成相同的最终结果？有谁知道如何做到这一点？

val data2 = spark.read.format("csv").option("header", 
"true").load("/data/c7.csv")
val goodBadRecords = data2.map(
  row =>{ 
  val n0 = row(0).toString.toLowerCase().toDouble
  val n1 = row(1).toString.toLowerCase().toDouble
  val n2 = row(2).toString.toLowerCase().toDouble
  val n3 = row(3).toString.toLowerCase().toDouble  
  (n0, Vectors.dense(n1,n2,n3))    
 }
).toDF("label", "features")

谢谢

问候，

阿迪尔

Answer 1

一个VectorAssembler可以完成这项工作：

VectorAssembler是一个变换器，它将给定的列列表组合到一个向量列中。将原始特征组合成单个特征向量非常有用

根据您的代码，解决方案将如下所示：

val data2 = spark.read.format("csv")
  .option("header","true")
  .option("inferSchema", "true") //1
  .load("/data/c7.csv")

val fields = data2.schema.fieldNames

val assembler = new VectorAssembler()
  .setInputCols(fields.tail) //2
  .setOutputCol("features") //3

val goodBadRecords = assembler.transform(data2)
  .withColumn("label", col(fields(0))) //4
  .drop(fields:_*) //5

备注：

1. 输入数据需要模式，因为VectorAssembler只接受以下输入列类型：所有数字类型，布尔类型和矢量类型（same link）。你似乎有一个双打的csv，所以推断架构应该工作。但是，当然，将字符串数据转换为双精度数的任何其他方法也是可以的。
2. 使用除第一列之外的所有列作为VectorAssembler的输入
3. 命名VectorAssembler功能的结果列
4. 创建一个名为label的新列作为第一列的副本
5. 删除所有原始列。最后一步是可选的，因为学习算法通常只查看标签和特征列并忽略所有其他列