Как объединить два фрейма данных pyspark для операции Arraytype?

Question

У меня есть два кадра данных, A и B. У каждого есть столбец с именем 'names', а этот столбец - ArrayType(StringType()).

Теперь я хочу оставить соединение A и B при условии, что A['names'] и B['names'] имеют общие элементы.

Вот пример:

A:

+---------------+
|          names|
+---------------+
|['Mike','Jack']|
|      ['Peter']|
+---------------+

B

+---------------+
|          names|
+---------------+
|['John','Mike']| 
|           null|
+---------------+

после левого соединения у меня должно быть:

+---------------+---------------+
|        A_names|        B_names|
+---------------+---------------+
|['Mike','Jack']|['John','Mike']| 
|      ['Peter']|           null|
+---------------+---------------+

Answer 1

В вашем случае вам нужно explode значения- explode будет производить одну строку на значение в ваших массивах, а затем вы можете присоединиться к ним и уменьшить конечный результат до желаемого формата.

В примере кода я разбил имена и присоединился к фреймам данных во вновь созданном столбце (B_names). Наконец, результат будет сгруппирован по "names" для удаления созданных дубликатов.

Для функции группировки по агрегату можно использовать pyspark.sql.functions.first() с параметром ignorenulls, установленным в True.

from pyspark.sql.functions import col, explode, first

test_list = [['Mike', 'Jack']], [['Peter']]
test_df = spark.createDataFrame(test_list, ["names"])

test_list2 =  [["John","Mike"]],[["Kate"]]
test_df2 = spark.createDataFrame(test_list2, ["names"])

test_df2 = test_df2.select(
    col("names").alias("B_names"),
    explode("names").alias("single_names")
)

test_df.select(col("names").alias("A_names"), explode("names").alias("single_names"))\
    .join(test_df2, on="single_names", how="left" )\
    .groupBy("A_names").agg(first("B_names", ignorenulls=True).alias("B_names")).show()

Результат:

+------------+------------+
|     A_names|     B_names|
+------------+------------+
|[Mike, Jack]|[John, Mike]|
|     [Peter]|        null|
+------------+------------+