如何在 Scala（Spark 2.0）中将带有字符串的 DataFrame 转换为带有 Vectors 的 DataFrame

Posted 2023-04-15

【中文标题】如何在 Scala（Spark 2.0）中将带有字符串的 DataFrame 转换为带有 Vectors 的 DataFrame

【英文标题】：How to convert a DataFrame with String into a DataFrame with Vectors in Scala(Spark 2.0)

【发布时间】：2016-12-05 21:00:51

【问题描述】：

我有一个 DataFrame，其中包含一个名为 KFA 的列，其中包含一个两端带有尖括号的字符串。这个长字符串中有 4 个双精度值。我想将其转换为带有向量的 DataFrame。

这是 DataFrame 的第一个元素：

> dataFrame1.first()
res130: org.apache.spark.sql.Row = [[.00663 .00197 .29809 .0034]]

你能帮我把它转换成一个有 4 个双精度值的密集向量吗？

我试过这个命令

dataFrame1.select("KFA")
    .map((x=>x.mkString("").replace("]","").replace("[","").split(" ")))
    .rdd.map(x=>Vectors.dense(x(0).toDouble,x(1).toDouble,x(2).toDouble,x(3).toDouble,x(4).toDouble))

这看起来非常笨拙且难以阅读。你能建议任何其他的方法吗？

【问题讨论】：

如果你只是要拆分mkString的目的是什么？

@cricket_007 使用 mkString 因为我无法在 spark.sql.Row 上使用 .replace("]","")

你不需要创建一个字符串，我不认为。您可以从 Row 对象中getAs[Double]

【参考方案1】：

这是一个带有正则表达式的选项：

import org.apache.spark.mllib.linalg.Vector, Vectors

val p = "[.0-9]+".r 
val rddVec = dataFrame1.select("KFA")
             .map(x => Vectors.dense(p.findAllIn(x(0).toString).map(_.toDouble).toArray))

# rddVec: org.apache.spark.rdd.RDD[org.apache.spark.mllib.linalg.Vector] = MapPartitionsRDD[49] at map at <console>:39

rddVec.collect
# res43: Array[org.apache.spark.mllib.linalg.Vector] = 
         Array([0.00663,0.00197,0.29809,0.0034], [0.00663,0.00197,0.29809,0.0034])