Comparing the values (list) of Pyspark dataframes

Question

I would like to make a comparison between two df1 df2 dataframes on the list_id column :

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

The desired result is:

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

My aim is to concatenate lists where their intersection is not empty and to keep the others as is with pyspark.

Note: my dataframes are very large, the use of a join with Spark Sql is impossible.

what if there are multiple matches between the two dataframes? — vb_rises
I tried to make a full join with a function of intersection between the columns of the lists, which generates a memory error — data

Xavier Canton Xavier Canton · Accepted Answer · 2019-10-01T12:46:41

I came up with a code which works without any join operation. It is somehow a pretty messy and I don't know how it will behave memory wize considering I'm exploding the array multiple times.

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

Comparing the values (list) of Pyspark dataframes

1 Answers