转换Spark DataFrame以进行ML处理

时间:2018-08-31 14:24:55

标签: scala apache-spark apache-spark-sql apache-spark-mllib apache-spark-ml

我编写了以下代码,将数据馈送到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")

谢谢

此致

阿黛尔

1 个答案:

答案 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

备注:

  1. 输入数据必须有一个架构,因为VectorAssembler仅接受以下输入列类型:所有数字类型,布尔类型和向量类型(same link)。您似乎有一个带有double的csv,因此推断该模式应该有效。但是,当然,将字符串数据转换为双精度的任何其他方法也是可以的。
  2. 使用第一列以外的所有内容作为VectorAssembler的输入
  3. 命名VectorAssembler的结果列功能
  4. 创建一个名为 label 的新列作为第一列的副本
  5. 删除所有原始列。最后一步是可选步骤,因为学习算法通常只查看label和feature列,而忽略所有其他列