scala - Spark 斯卡拉 : convert arbitrary N columns into Map

Question

我有以下数据结构，表示电影 ID(第一列)和其余列中不同用户对该电影的评分 - 类似这样:

+-------+----+----+----+----+----+----+----+----+----+----+----+----+----+----+----+
|movieId|   1|   2|   3|   4|   5|   6|   7|   8|   9|  10|  11|  12|  13|  14|  15|
+-------+----+----+----+----+----+----+----+----+----+----+----+----+----+----+----+
|   1580|null|null| 3.5| 5.0|null|null|null|null|null|null|null|null|null|null|null|
|   3175|null|null|null|null|null|null|null|null|null|null|null|null|null| 5.0|null|
|   3794|null|null|null|null|null|null|null|null|null|null|null| 3.0|null|null|null|
|   2659|null|null|null| 3.0|null|null|null|null|null|null|null|null|null|null|null|

我想把这个DataFrame转换成一个DataSet

最终案例类 MovieRatings(movie_id: Long, ratings: Map[Long, Double])

所以它会是这样的

[1580, [1 -> null, 2 -> null, 3 -> 3.5, 4 -> 5.0, 5 -> null, 6 -> null, 7 -> null,...]]

等等

如何做到这一点？

这里的问题是用户数量是任意的。我想将它们压缩到一个列中，而第一列保持不变。

Answer 1

首先，您必须将您的 DataFrame 转换为具有与您的案例类匹配的架构的 DataFrame，然后您可以使用 .as[MovieRatings] 将 DataFrame 转换为 Dataset[MovieRatings]:

import org.apache.spark.sql.functions._
import spark.implicits._

// define a new MapType column using `functions.map`, passing a flattened-list of
// column name (as a Long column) and column value
val mapColumn: Column = map(df.columns.tail.flatMap(name => Seq(lit(name.toLong), $"$name")): _*)

// select movie id and map column with names matching the case class, and convert to Dataset:
df.select($"movieId" as "movie_id", mapColumn as "ratings")
  .as[MovieRatings]
  .show(false)