将 pyspark dataframe 数组列中每个数组的值替换为其对应的 id

Question

I have a pyspark.sql dataframe that looks like this:

id	name	refs
1	A	B, C,D
2	B	A
3	C	A, B

I'm trying to build a function that replaces the values of each array in ref by the corresponding ID of the name that it references and if there's no matching name in the Name column then it would ideally filter that value out or set it to null . The results would ideally look something like this:

id	name	refs
1	A	2, 3
2	B	1
3	C	1, 2

I tried doing this by defining a UDF that collects all names from the table and then obtains the indices of the intersection between each ref array and the set of all names. It works but is extremely slow, I'm sure there's probably better ways to do this using Spark and/or SQL.

Answer 1

You can explode the arrays, do a self-join using the exploded ref and name, and collect the joined ids back to an array using collect_list .

import pyspark.sql.functions as F

df1 = df.select('id', 'name', F.explode('refs').alias('refs'))
df2 = df.toDF('id2', 'name2', 'refs2')

result = df1.join(df2, df1.refs == df2.name2) \
            .select('id', 'name', 'id2') \
            .groupBy('id', 'name') \
            .agg(F.collect_list('id2').alias('refs'))

result.show()
+---+----+------+
| id|name|  refs|
+---+----+------+
|  1|   A|[2, 3]|
|  2|   B|   [1]|
|  3|   C|[1, 2]|
+---+----+------+