比较Pyspark数据框的值(列表)

时间:2019-10-01 09:56:24

标签: python pyspark pyspark-dataframes

我想在list_id列上的两个df1 df2数据帧之间进行比较:

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 :(得分:0)

我想出了一个无需任何连接操作即可工作的代码。 考虑到我多次分解数组,这有点混乱,我也不知道它将如何表现内存浪费。

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()
    )