获取spark数据帧中二维包裹的每个元素的第一个元素

时间:2018-05-20 09:04:39

标签: scala apache-spark apache-spark-sql

我按https://spark.apache.org/docs/2.3.0/api/scala/index.html#org.apache.spark.ml.recommendation.ALSModel recommendForUserSubset方法获得了一个数据框,

val recsysResult = model.recommendForAllUsers(5)
recsysResult.show(2,false)

+-------+------------------------------------------------------------------------------------------+
|user_id|recommendations                                                                           |
+-------+------------------------------------------------------------------------------------------+
|1      |[[111, 0.9899166], [110, 0.9899166], [101, 0.9899166], [100, 0.9899166], [102, 0.9899166]]|
|0      |[[0, 0.9899168], [10, 0.9899168], [30, 0.9899168], [40, 0.9899168], [50, 0.9899168]]      |
+-------+------------------------------------------------------------------------------------------+  

现在我想获得一个像

这样的新数据框
1 | 111 | 110 | 101 | 100 | 102  
0 | 0   | 10  | 30  | 40  | 50 

获取二维包裹阵列recommendations的每个元素的第一个元素。我该怎么办?

1 个答案:

答案 0 :(得分:0)

您拥有的结构不是嵌套数组。它是一个结构数组,所以要从中提取数据,而不使用udf,我们必须知道这些名称。

由于您省略了架构,我们必须首先重命名对象:

val n: Int = 5
val recsysResult = model.recommendForAllUsers(n)

val  recsysResultIds = recsysResult.withColumn(
  "ids",
  col("recommendations").cast("array<struct<id:int,rating:float>>")("id")
)

然后选择

recsysResultIds.select($"user_id" +: (0 to n).map(col("ids")(_)): _* )