Объединить несколько столбцов из одного фрейма данных в один столбец из другого без операции множественного объединения в pyspark

Question

Я пытаюсь сопоставить несколько столбцов из одного фрейма данных (df) с многоязычным словарем (df_label) и извлечь соответствующие метки для каждого столбца.
Примечание. Это не двойной вопрос Соединение нескольких столбцов из одной таблицы в один столбец из другой таблицы

Ниже приведен пример фреймов данных df и df_label и желаемого результата

  df                df_label                       output
+---+---+       +---+-----+----+        +---+---+------+------+------+
|  s|  o|       |  e| name|lang|        |  s|  o|s_name|o_name|  lang|
+---+---+       +---+-----+----+        +---+---+------+------+------+
| s1| o1|       | s1|s1_en|  en|        | s2| o1| s2_fr| o1_fr|    fr|
| s1| o3|       | s1|s1_fr|  fr|        | s1| o1| s1_fr| o1_fr|    fr|
| s2| o1|       | s2|s2_fr|  fr|        | s1| o1| s1_en| o1_en|    en|
| s2| o2|       | o1|o1_fr|  fr|        | s2| o2| s2_fr| o2_fr|    fr|
+---+---+       | o1|o1_en|  en|        +---+---+------+------+------+
                | o2|o2_fr|  fr|
                +---+-----+----+

Другими словами, я хочу сопоставить оба столбца [s, o] из df со столбцом e из df_label и найти их соответствующие метки на разных языках, как показано выше.

Многоязычный словарь (df_label) огромен, а столбцы [s, o] имеют много дубликатов, поэтому две операции соединения крайне неэффективны.

Есть ли какой-нибудь способ, который может быть достигнут без нескольких объединений?

К вашему сведению, это то, что я делал, используя несколько объединений, но мне действительно это не нравится.

df = spark.createDataFrame([('s1','o1'),('s1','o3'),('s2','o1'),('s2','o2')]).toDF('s','o')
df_label = spark.createDataFrame([('s1','s1_en','en'),('s1','s1_fr','fr'),('s2','s2_fr','fr'),('o1','o1_fr','fr'),('o1','o1_en','en'),('o2','o2_fr','fr')]).toDF('e','name','lang')
df = df.join(df_label,col('s')==col('e')).drop('e').withColumnRenamed('name','s_name').withColumnRenamed('lang','s_lang')
df = df.join(df_label,col('o')==col('e')).drop('e').withColumnRenamed('name','o_name').select('s','o','s_name','o_name','s_lang','o','o_name','lang').withColumnRenamed('lang','o_lang').filter(col('o_lang')==col('s_lang')).drop('s_lang')   

Answer 1

Опираясь на то, что gaw предложил, это мое предлагаемое решение
Подход состоял в том, чтобы использовать только одно объединение, а затем использовать условный агрегат collect_list, чтобы проверить, было ли совпадение для столбца s или o .

df = = spark.createDataFrame([('s1','o1'),('s1','o3'),('s2','o1'),('s2','o2')]).toDF('s','o')
df_label = spark.createDataFrame([('s1','s1_en','en'),('s1','s1_fr','fr'),('s2','s2_fr','fr'),('o1','o1_fr','fr'),('o1','o1_en','en'),('o2','o2_fr','fr')]).toDF('e','name','lang')

df.join(df_label,(col('e')== col('s')) | (col('e') == col('o'))) \
.groupBy(['s','o','lang']) \
.agg(collect_list(when(col('e')==col('s'),col('name'))).alias('s_name')\
,collect_list(when(col('e')==col('o'),col('name'))).alias('o_name')) \
.withColumn('s_name',explode('s_name')).withColumn('o_name',explode('o_name')).show()

    +---+---+----+------+------+
    |  s|  o|lang|s_name|o_name|
    +---+---+----+------+------+
    | s2| o2|  fr| s2_fr| o2_fr|
    | s1| o1|  en| s1_en| o1_en|
    | s1| o1|  fr| s1_fr| o1_fr|
    | s2| o1|  fr| s2_fr| o1_fr|
    +---+---+----+------+------+

Answer 2

Я создал способ, который работает только с одним объединением, но так как он использует дополнительные (дорогие) операции, такие как explode и т. Д. Я не уверен, что это быстрее. Но если вам нравится, вы можете попробовать.

Следующий код выдает желаемый результат:

df = spark.createDataFrame([('s1','o1'),('s1','o3'),('s2','o1'),('s2','o2')]).toDF('s','o')
df_label = spark.createDataFrame([('s1','s1_en','en'),('s1','s1_fr','fr'),('s2','s2_fr','fr'),('o1','o1_fr','fr'),('o1','o1_en','en'),('o2','o2_fr','fr')]).toDF('e','name','lang')
df = df.join(df_label,[(col('s')==col('e')) | \
  (col('o')==col('e'))]).drop('e').\   #combine the two join conditions
  withColumn("o_name",when(col("name").startswith("o"),col("name")).otherwise(None)).\
  withColumn("s_name",when(col("name").startswith("s"),col("name")).otherwise(None)).\ #create the o_name and s_name cols
  groupBy("s","o").agg(collect_list("o_name").alias("o_name"),collect_list("s_name").alias("s_name")).\
  #perform a group to aggregate the required vales
  select("s","o",explode("o_name").alias("o_name"),"s_name").\ # explode the lists from the group to attach it to the correct pairs of o and s
  select("s","o",explode("s_name").alias("s_name"),"o_name").\
  withColumn("o_lang", col("o_name").substr(-2,2)).\
  withColumn("lang", col("s_name").substr(-2,2)).filter(col("o_lang")==col("lang")).drop("o_lang")
  #manually create the o_lang and lang columns 

Результат:

+---+---+------+------+----+
|s  |o  |s_name|o_name|lang|
+---+---+------+------+----+
|s2 |o2 |s2_fr |o2_fr |fr  |
|s2 |o1 |s2_fr |o1_fr |fr  |
|s1 |o1 |s1_fr |o1_fr |fr  |
|s1 |o1 |s1_en |o1_en |en  |
+---+---+------+------+----+