Pyspark 基于另一个 dataframe 替换数组列上的值

Question

我有两个数据框，一个只是带有一些具有关联names的唯一ids ，如下所示：

Id name
0 name_a
1 name_b
2 name_c

第二个 dataframe 包含存储在数组中的第一个 dataframe 的 id，在每一行中：

Row_1 row_2
0 [0,2]
1 [1,0]

我的问题是可以从第二个 dataframe 替换 arrays ，因此它会根据ids检查第一个 df 的names ，所以：

Row_1 row_2
0 [name_a, name_c]
1 [name_b, name_a]

创建第一个 df 的 map 并使用 udf 将其添加到第二个 df 似乎太耗时了。 非常感谢任何有关如何解决此问题的帮助。

Answer 1

您可以尝试使用explode function 将数组转换为行，然后将数据与初始数据框连接起来，在最后一步通过 &.agg(collect_list()) 进行分组

from pyspark.sql.functions import explode

df3 = df2.select(df2.row_1,explode(df2.row_2))
df4 = df3.join(df1,df3.row_1==df1.Id).select(df3.row_1,df1.name)
df5 = df4.groupBy('row_1').agg(collect_list('name').alias('name'))

参考链接：

1. https://sparkbyexamples.com/pyspark/pyspark-explode-array-and-map-columns-to-rows/#:~:text=explode%20%E2%80%93%20PySpark%20explode%20array%20或，它%20包含%20all%20array%20elements 。
2. https://www.owenrumney.co.uk/pyspark-opposite-of-explode/

Answer 2

使用array_contains function + groupby 和collect_list加入：

from pyspark.sql import functions as F

df1 = spark.createDataFrame([(0, "name_a"), (1, "name_b"), (2, "name_c")], ["Id", "name"])
df2 = spark.createDataFrame([(0, [0, 2]), (1, [1, 0])], ["Row_1", "Row_2"])

result = df2.join(
    df1, on=F.array_contains("Row_2", F.col("Id")), how="left"
).groupBy("Row_1").agg(
    F.collect_list("name").alias("Row_2")
)

result.show()
#+-----+----------------+
#|Row_1|           Row_2|
#+-----+----------------+
#|    0|[name_a, name_c]|
#|    1|[name_a, name_b]|
#+-----+----------------+