Scala Spark-仅对几个最小的项目进行分组和聚合

简言之

我有两个数据帧和函数的笛卡尔积（交叉连接），它为这个积的给定元素提供了一些分数。现在，我想为第一个DF的每个成员获取第二个DF的几个“最佳匹配”元素

详细信息

下面是一个简化的示例，因为我的实际代码中有一些额外的字段和过滤器

给定两组数据，每组都有一些id和值：

// simple rdds of tuples
val rdd1 = sc.parallelize(Seq(("a", 31),("b", 41),("c", 59),("d", 26),("e",53),("f",58)))
val rdd2 = sc.parallelize(Seq(("z", 16),("y", 18),("x",3),("w",39),("v",98), ("u", 88)))

// convert them to dataframes:
val df1 = spark.createDataFrame(rdd1).toDF("id1", "val1")
val df2 = spark.createDataFrame(rdd2).toDF("id2", "val2")

对于来自第一个和第二个数据集的一对元素，某些函数给出它们的“匹配分数”：

我们可以创建两组的乘积，并计算每对的分数：

val dfc = df1.crossJoin(df2)
val r = dfc.withColumn("rez", fu(col("val1"), col("val2")))
r.show

+---+----+---+----+---+
|id1|val1|id2|val2|rez|
+---+----+---+----+---+
|  a|  31|  z|  16|  8|
|  a|  31|  y|  18| 10|
|  a|  31|  x|   3|  2|
|  a|  31|  w|  39| 15|
|  a|  31|  v|  98| 13|
|  a|  31|  u|  88|  2|
|  b|  41|  z|  16| 14|
|  c|  59|  z|  16| 12|
...

现在我们想让这个结果按

id1

分组：

r.groupBy("id1").agg(collect_set(struct("id2", "rez")).as("matches")).show

+---+--------------------+
|id1|             matches|
+---+--------------------+
|  f|[[v,2], [u,8], [y...|
|  e|[[y,5], [z,3], [x...|
|  d|[[w,2], [x,6], [v...|
|  c|[[w,2], [x,6], [v...|
|  b|[[v,2], [u,8], [y...|
|  a|[[x,2], [y,10], [...|
+---+--------------------+

但实际上，我们只想保留少数（比如3场）的“比赛”，即得分最高（比如最低）的比赛

问题是

如何将“匹配项”排序并缩减为前N个元素？可能是关于collect_list和sort_数组的，尽管我不知道如何按内部字段排序

是否有办法确保在大输入DFs情况下进行优化-例如，在聚合时直接选择最小值。我知道如果我在编写代码时不使用火花-为每个

id1

保留小数组或优先级队列，并在应该添加的地方添加元素，可能会删除以前添加的一些元素，那么就很容易做到这一点

---

交叉连接是一项代价高昂的操作，这没关系，但我想避免在结果上浪费内存，而在下一步，我将放弃大部分结果。我的实际用例处理的DFs条目少于1 mln，因此交叉连接仍然可行，但由于我们只希望为每个

id1

选择10-20个最匹配项，因此似乎不希望在步骤之间保留不必要的数据

对于start，我们只需要取前n行。为此，我们用“id1”对DF进行分区，并用res对组进行排序。我们使用它向DF添加行数列，就像我们可以使用where函数获取前n行一样。这样你就可以继续写你写的代码了。按“id1”分组并收集列表。只是现在您已经拥有了最高的行

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

val n = 3
val w = Window.partitionBy($"id1").orderBy($"res".desc)
val res = r.withColumn("rn", row_number.over(w)).where($"rn" <= n).groupBy("id1").agg(collect_set(struct("id2", "res")).as("matches"))

---

在这里，我们创建一个udf，它接受数组列和整数值n。UDF用您的“RES”对数组进行排序，只返回第一个n个元素。

也许考虑使用窗口函数<代码> Reals<代码>来实现这一点。或者编写一个udf，从您的collect_集合中获取结果数组的前3名。谢谢！我需要一些时间来研究你的解释-我会回来的@RodionGorkovenko不客气。你能批准这个答案，让人们看到你的问题已经得到了回答吗？格雷西亚斯。

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

val n = 3
val w = Window.partitionBy($"id1").orderBy($"res".desc)
val res = r.withColumn("rn", row_number.over(w)).where($"rn" <= n).groupBy("id1").agg(collect_set(struct("id2", "res")).as("matches"))

val sortTakeUDF = udf{(xs: Seq[Row], n: Int)} => xs.sortBy(_.getAs[Int]("res")).reverse.take(n).map{case Row(x: String, y:Int)}}
r.groupBy("id1").agg(sortTakeUDF(collect_set(struct("id2", "res")), lit(n)).as("matches"))