Spark DataFrame中提取二维WrappedArray各元素首个元素的方法

我来帮你搞定这个需求！你现在需要把recommendForAllUsers返回的DataFrame里，recommendations数组中的每个元素的商品ID提取出来，转换成一行一个用户、后面跟着多个推荐商品ID的宽表结构对吧？

实现思路 #

recommendations列是一个包含结构体的数组，每个结构体里存着推荐的商品ID和对应的评分。我们可以通过数组索引定位到每个推荐项，再提取其中的商品ID字段，最后生成新的列即可。

具体代码实现 #

首先先确认你的recommendations列里的结构体字段名，你可以先执行recsysResult.printSchema()查看，比如默认情况下ALS返回的结构体字段是product（商品ID）和rating（评分），如果你的字段名是item_id就替换成对应的名称。

下面是针对你示例中5个推荐项的Scala代码：

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

// 提取每个推荐位置的商品ID，生成新列
val wideResult = recsysResult
  // 提取第1个推荐的商品ID
  .withColumn("rec_1", col("recommendations")(0)("product"))
  // 提取第2个推荐的商品ID
  .withColumn("rec_2", col("recommendations")(1)("product"))
  // 提取第3个推荐的商品ID
  .withColumn("rec_3", col("recommendations")(2)("product"))
  // 提取第4个推荐的商品ID
  .withColumn("rec_4", col("recommendations")(3)("product"))
  // 提取第5个推荐的商品ID
  .withColumn("rec_5", col("recommendations")(4)("product"))
  // 删除原来的recommendations数组列
  .drop("recommendations")

// 查看转换后的结果
wideResult.show(false)

代码说明 #

• col("recommendations")(index)：通过索引获取数组中第index个推荐项（索引从0开始）
• ("product")：从推荐项的结构体中提取商品ID字段，如果你实际的字段名是item_id，替换成("item_id")即可
• withColumn：依次添加5个推荐商品ID的列，你可以根据需要修改列名（比如直接叫1、2...但列名最好符合规范）
• drop("recommendations")：移除原来的数组列，得到你想要的宽表结构

执行后得到的结果就会和你期望的一致啦：

+-------+-----+-----+-----+-----+-----+
|user_id|rec_1|rec_2|rec_3|rec_4|rec_5|
+-------+-----+-----+-----+-----+-----+
|1      |111  |110  |101  |100  |102  |
|0      |0    |10   |30   |40   |50   |
+-------+-----+-----+-----+-----+-----+

内容的提问来源于stack exchange，提问作者superDuck