如何根据其他数据框的两列向数据框添加新列

Question

我有两个 DataFrame df_data 和 df_node_labels：

df_data =

    nodeId   field1
    1        abc
    2        def
    3        fed
    4        kfl

df_node_labels =
    srcId   srcLabel    dstId    dstLabel
    1       AAA         2        BBB
    2       BBB         4        FFF
    4       FFF         3        CCC

我想在df_data 中添加一列label。 label的值应该取自srcLabel和dstLabel：

这就是我尝试获取标签信息的方式：

var df = df_data.join(df_node_labels.select("srcId","srcLabel"),col("nodeId")===col("srcId"),"left")
df = df.join(df_node_labels.select("dstId","dstLabel"),col("nodeId")===col("dstId"),"left")

但是，这会在 df 中创建两列 srcLabel 和 dstLabel，而我只想得到一列 label。

这是预期的结果：

df =
        nodeId   field1   label
        1        abc      AAA
        2        def      BBB
        3        fed      CCC
        4        kfl      FFF

更新：

我可以这样做，但在我看来，做一件简单的事情还有很长的路要走：

df = df.withColumn("label", when(col("srcLabel") =!= "", col("srcLabel")).otherwise(col("dstLabel"))).drop("srcLabel").drop("dstLabel")

Answer 1

您可以从df_node_labels 创建一个唯一数据作为下面的finalDF 并执行join 操作，这将为您提供预期的结果。

val finalDF = df_node_labels.select($"srcId".as("nodeId"), $"srcLabel".as("label"))
  .union(
    df_node_labels.select($"dstId".as("nodeId"), $"dstLabel".as("label"))
  ).dropDuplicates()


df_data.join(finalDF, Seq("nodeId"), "left")
  .show(false)

输出：

+------+------+-----+
|nodeId|field1|label|
+------+------+-----+
|1     |abc   |AAA  |
|2     |def   |BBB  |
|3     |fed   |CCC  |
|4     |kfl   |FFF  |
+------+------+-----+

希望对你有所帮助！