Spark Scala:将任意N列转换为Map

时间:2017-10-19 14:04:59

标签: scala apache-spark dataframe dataset

我有以下数据结构代表电影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

final case class 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,...]]

如何做到这一点?

这里的事情是用户数量是任意的。我希望将它们压缩成一列,不要触及第一列。

2 个答案:

答案 0 :(得分:3)

首先,您必须将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)

答案 1 :(得分:3)

您可以使用spark.sql.functions.map从任意列创建地图。它期望键和值之间的序列交替,可以是列类型或字符串。这是一个例子:

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

case class Input(movieId: Int, a: Option[Double], b: Option[Double], c: Option[Double])

val data = Input(1, None, Option(3.5), Option(1.4)) :: 
        Input(2, Option(4.2), Option(1.34), None) :: 
        Input(3, Option(1.11), None, Option(3.32)) :: Nil

val df = sc.parallelize(data).toDF

// Exclude the PK column from the map
val mapKeys = df.columns.filterNot(_ == "movieId")

// Build the sequence of key, value, key, value, ..
val pairs = mapKeys.map(k => Seq(lit(k), col(k))).flatten

val mapped = df.select($"movieId", functions.map(pairs:_*) as "map")
mapped.show(false) 

生成此输出:

+-------+------------------------------------+
|movieId|map                                 |
+-------+------------------------------------+
|1      |Map(a -> null, b -> 3.5, c -> 1.4)  |
|2      |Map(a -> 4.2, b -> 1.34, c -> null) |
|3      |Map(a -> 1.11, b -> null, c -> 3.32)|
+-------+------------------------------------+