Сравнение значений (списка) фреймов данных Pyspark - PullRequest
Новая
       11

Сравнение значений (списка) фреймов данных Pyspark

0 голосов
/ 01 октября 2019

Я хотел бы сравнить два кадра данных df1 df2 в столбце list_id: 

df1 = 
+---------+
|  list_id|
+---------+
|[1, 2, 3]|
|[4, 5, 6]|
|[7, 8, 9]|
+---------+
df2 =
+------------+
|     list_id|
+------------+
| [10, 3, 11]|
|[12, 13, 14]|
| [15, 6, 16]|
+------------+

Желаемый результат: 

df2 =
+-------------------+
|            list_id|
+-------------------+
| [1, 2, 3, 10, 11] |
| [4, 5, 6, 15, 16] |
| [7, 8, 9]         |
| [12, 13, 14]      |
+-------------------+

Моя цель - объединить списки, в которыхих пересечение не пустое, и остальные остаются такими же, как с pyspark.

Примечание: мои кадры данных очень велики, использование объединения с Spark Sql невозможно.

1 Ответ

0 голосов
/ 01 октября 2019

Я придумал код, который работает без какой-либо операции соединения. Это как-то довольно грязно, и я не знаю, как это будет вести себя в отношении памяти, учитывая, что я взрываю массив несколько раз. 

import pyspark.sql.functions as F
from pyspark.sql.window import Window

df1 = (sc.parallelize([(1, 2, 3), (4, 5, 6), (7, 8, 9)])
         .toDF(('c1', 'c2', 'c3'))
         .select(F.array(F.col('c1'), F.col('c2'), F.col('c3')).alias('id_list'))
        )

df2 = (sc.parallelize([(10, 3, 11), (12, 13, 14), (15, 6, 16)])
         .toDF(('c1', 'c2', 'c3'))
         .select(F.array(F.col('c1'), F.col('c2'), F.col('c3')).alias('id_list'))
         )

out = (df1.union(df2)
         .withColumn('key1', F.explode('id_list'))
         .withColumn('key2', F.explode('id_list'))
         .groupBy('key1')
         .agg(F.sort_array(F.collect_set(F.col('key2'))).alias('id_list'))
         .withColumn('key1', F.explode('id_list'))
         .withColumn('max_length', F.max(F.size('id_list')).over(Window().partitionBy('key1')))
         .where(F.col('max_length')==F.size('id_list'))
         .select('id_list')
         .distinct()
    )
...