Как добавить новый столбец во фрейм данных на основе двух столбцов других фреймов данных

Question

У меня есть два кадра данных 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

Я хочу добавить столбец label к df_data. Значения 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")

Однако это создает два столбца srcLabel и dstLabel в df, в то время как я хочу получить только один столбец 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  |
+------+------+-----+

Надеюсь, это вам помогло!