从 DataFrame 中收集两个值，并将它们用作案例类的参数；寻找不那么冗长的解决方案

Posted 2023-04-15

【中文标题】从 DataFrame 中收集两个值，并将它们用作案例类的参数；寻找不那么冗长的解决方案

【英文标题】：Collecting two values from a DataFrame, and using them as parameters for a case class; looking for less verbose solution

【发布时间】：2019-03-09 13:18:37

【问题描述】：

我在 spark 中有一些数据，result: DataFrame = ...，其中两个整数列是感兴趣的； week 和 year。这些列的值对于所有行都是相同的。

我想提取这两个整数值，并将它们作为参数传递以创建WeekYear：

case class WeekYear(week: Int, year: Int)

以下是我目前的解决方案，但我认为必须有一种更优雅的方式来做到这一点。如果没有创建temp 的中间步骤，如何做到这一点？

val temp = result
  .select("week", "year")
  .first
  .toSeq
  .map(_.toString.toInt)

val resultWeekYear = WeekYear(temp(0), temp(1))

【参考方案1】：

将case class 与数据框结合使用的最佳方法是允许spark 使用.as() 方法将其转换为数据集。只要您的案例类具有与所有列名匹配的属性，它就应该很容易工作。

case class WeekYear(week: Int, year: Int)
val df = spark.createDataset(Seq((1, 1), (2, 2), (3, 3))).toDF("week", "year")
val ds = df.as[WeekYear]
ds.show()

它提供了一个看起来像这样的Dataset[WeekYear]：

+----+----+
|week|year|
+----+----+
|   1|   1|
|   2|   2|
|   3|   3|
+----+----+

您可以使用一些更复杂的嵌套类，但您必须开始使用Encoders，以便 spark 知道如何来回转换。

Spark 会进行一些隐式转换，因此ds 可能仍然看起来像Dataframe，但它实际上是强类型Dataset[WeekYear]，而不是具有任意列的Dataset[Row] .您对它的操作类似于RDD。然后只需抓住其中一个.first()，您就已经拥有了您需要的类型。

val resultWeekYear = ds.first