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

时间:2017-12-15 19:59:06

标签: scala apache-spark spark-dataframe

我有一个火花数据框,其中包含每行的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查找但如果第二个表很大,我怀疑这将非常慢。

2 个答案:

答案 0 :(得分:3)

您可以首先展开您的ID列表,然后在其ID上加入两个Dataframe,最后收集列表中的说明。

我们假设df是您的第一个数据框,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"))

答案 1 :(得分:1)

如果您的第二个表很大且第一个表相对较小,您可以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]|
+-----+------------+