将列表过滤到scala数据集中每个参数值的前2个案例类

Question

我有一个像这样的spark数据集：

+--------+--------------------+
|     uid|     recommendations|
+--------+--------------------+
|41344966|[[2133, red]...|
|41345063|[[11353, red...|
|41346177|[[2996, yellow]...|
|41349171|[[8477, green]...|

res98: org.apache.spark.sql.Dataset[userItems] = [uid: int, recommendations: array<struct<iid:int,color:string>>]

我想过滤每个推荐数组，以包含每种颜色的前两个。伪示例：

[(13,'red'), (4,'green'), (8,'red'), (2,'red'), (10, 'yellow')]

会成为

[(13,'red'), (4,'green'), (8,'red'), (10, 'yellow')]

如何使用数据集在scala中有效地执行此操作？有没有像reduceGroups这样的优雅解决方案？

到目前为止我所拥有的：

case class itemData (iid: Int, color: String)

val filterList = (recs: Array[itemData], filterAttribute, maxCount) => {
  // filter the list somehow... using the max count and attribute
  })

dataset.map(d => filterList(d.recommendations, "color", 2))

Answer 1

您可以分解建议，然后通过uid和颜色创建行号分区，最后过滤掉大于2的行号。代码应如下所示。我希望它有所帮助。

//Creating Test Data
val df = Seq((13,"red"), (4,"green"), (8,"red"), (2,"red"), (10, "yellow")).toDF("iid", "color")
  .withColumn("uid", lit(41344966))
  .groupBy("uid").agg(collect_list(struct("iid", "color")).as("recommendations"))

df.show(false)
+--------+----------------------------------------------------+
|uid     |recommendations                                     |
+--------+----------------------------------------------------+
|41344966|[[13,red], [4,green], [8,red], [2,red], [10,yellow]]|
+--------+----------------------------------------------------+

val filterDF = df.withColumn("rec", explode(col("recommendations")))
    .withColumn("iid", col("rec.iid"))
    .withColumn("color", col("rec.color"))
    .drop("recommendations", "rec")
    .withColumn("rownum",
      row_number().over(Window.partitionBy("uid", "color").orderBy(col("iid").desc)))
    .filter(col("rownum") <= 2)
    .groupBy("uid").agg(collect_list(struct("iid", "color")).as("recommendations"))

filterDF.show(false)
+--------+-------------------------------------------+
|uid     |recommendations                            |
+--------+-------------------------------------------+
|41344966|[[4,green], [13,red], [8,red], [10,yellow]]|
+--------+-------------------------------------------+