Spark数据帧：使用第二个数据帧查找数组的元素

Question

我有一个spark数据框，其中包含每行的ID列表：

ident  list_of_ids
1      [3,4,5]
2      [5,6]
3      [2]
4      []

第二个将id映射到某些文本描述的spark数据帧：

id     desc
2      "aa"
3      "bb"
4      "cc"
5      "dd"
6      "ee"

是否有一种简单（快速）的方法将第三列添加到第一个数据框，其中包含与列表中的ID对应的描述列表：

ident  list_of_ids  list_of_desc
1      [3,4,5]      ["bb", "cc", "dd"]
2      [5,6]        ["dd", "ee"]
3      [2]          ["aa"]
4      []           []

我知道我可以在第二列上执行forEach并根据每个id进行查找，但如果第二个表很大，我怀疑这将非常慢。

Answer 1

您可以先爆炸您的ID列表，然后在其ID上加入两个Dataframe，最后收集列表中的描述。

我们假设df是您的第一个Dataframe，dfDesc包含以下描述：

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

df.withColumn("id", explode($"list_of_ids"))
  .join(dfDesc, Seq("id"))
  .groupBy($"ident", $"list_of_ids").agg(collect_list($"desc"))

Answer 2

如果你的第二个表很大而且第一个表相对较小，你可以1）从第一个表创建一个不同id的数据帧，并使用broadcast和第二个表连接，用desc-mapped列创建一个相对较小的数据帧，和2 ）加入两个创建的数据帧并使用groupBy / collect_list生成结果：

val df1 = Seq(
  (1, Seq(3, 4, 5)),
  (2, Seq(5, 6)),
  (3, Seq(2)),
  (4, Seq())
).toDF("ident", "list_of_ids")

val df2 = Seq(
  (2, "aa"),
  (3, "bb"),
  (4, "cc"),
  (5, "dd"),
  (6, "ee")
).toDF("id", "desc")

val df1Exploded = df1.select($"ident", explode($"list_of_ids").as("id"))

val df1Distinct = df2.join(broadcast(df1Exploded.select($"id").distinct), Seq("id"))

val dfResult = df1Exploded.join(df1Distinct, Seq("id")).groupBy($"ident").
  agg(collect_list($"desc").as("list_of_desc"))

dfResult.show
+-----+------------+
|ident|list_of_desc|
+-----+------------+
|    1|[bb, cc, dd]|
|    3|        [aa]|
|    2|    [dd, ee]|
+-----+------------+