我编写了以下代码,将数据馈送到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")
谢谢
此致
阿黛尔
答案 0 :(得分: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
备注: