scala - 如何将DataFrame中的struct映射到case类？

Question

在我的应用程序中的某个时刻，我有一个 DataFrame 和一个从 case 类创建的 Struct 字段。现在我想将它转换/映射回 case 类类型:

import spark.implicits._
case class Location(lat: Double, lon: Double)

scala> Seq((10, Location(35, 25)), (20, Location(45, 35))).toDF
res25: org.apache.spark.sql.DataFrame = [_1: int, _2: struct<lat: double, lon: double>]

scala> res25.printSchema
root
 |-- _1: integer (nullable = false)
 |-- _2: struct (nullable = true)
 |    |-- lat: double (nullable = false)
 |    |-- lon: double (nullable = false)

和基本的:

res25.map(r => {
   Location(r.getStruct(1).getDouble(0), r.getStruct(1).getDouble(1))
}).show(1)

看起来真的很脏
有没有更简单的方法？

Answer 1

在 Spark 1.6+ 中，如果要保留保留的类型信息，则使用 Dataset (DS)，而不是 DataFrame (DF)。

import spark.implicits._
case class Location(lat: Double, lon: Double)

scala> Seq((10, Location(35, 25)), (20, Location(45, 35))).toDS
res25: org.apache.spark.sql.Dataset[(Int, Location)] = [_1: int, _2: struct<lat: double, lon: double>]

scala> res25.printSchema
root
 |-- _1: integer (nullable = false)
 |-- _2: struct (nullable = true)
 |    |-- lat: double (nullable = false)
 |    |-- lon: double (nullable = false)

它会给你 Dataset[(Int, Location)] .现在，如果你想再次回到它的案例类起源，那么只需这样做:

scala> res25.map(r => r._2).show(1)
+----+----+
| lat| lon|
+----+----+
|35.0|25.0|
+----+----+

但是，如果你想坚持使用 DataFrame API，由于它是动态类型的，那么你必须像这样编码:

scala> res25.select("_2.*").map(r => Location(r.getDouble(0), r.getDouble(1))).show(1)
+----+----+
| lat| lon|
+----+----+
|35.0|25.0|
+----+----+