使用 Spark 数据框映射 ArrayType 列中的值

Question

我有一个带有 ArrayType 列的 Spark 数据框：

|id|neighbors|
+--+---------+
|a |[b,c]    |
|b |[a,d]    |
|c |[a]      |
|d |[b]      |
+--+---------+

我需要将此 ArrayType 列中的值映射到初始数据框。
期望的输出：

|id|neighbors    |
+--+-------------+
|a |[[a,d],[a]]  |
|b |[[b,c],[b]]  |
|c |[[b,c]]      |
|d |[[a,d]]      |
+--+-------------+

处理这个问题的最佳方法是什么？我的数据量非常大（大约1亿条记录）。

Answer 1

您需要 explode 列 'neighbors'，然后 join。此外，由于这是自连接，建议在数据帧上使用 alias。

初始 df:

from pyspark.sql import functions as F

df = spark.createDataFrame(
    [('a', ['b', 'c']),
     ('b', ['a', 'd']),
     ('c', ['a']),
     ('d', ['b'])],
    ['id', 'neighbors']
)

转型：

df = (
    df.withColumn('_neighbors', F.explode('neighbors')).alias('df1')
    .join(df.alias('df2'), F.col('df1._neighbors') == F.col('df2.id'))
    .groupBy('df1.id')
    .agg(F.collect_list('df2.neighbors').alias('neighbors'))
)
df.show()
# +---+-------------+
# | id|    neighbors|
# +---+-------------+
# |  d|     [[a, d]]|
# |  c|     [[b, c]]|
# |  b|[[b, c], [b]]|
# |  a|[[a, d], [a]]|
# +---+-------------+

使用 Spark 数据框映射 ArrayType 列中的值

Map values in ArrayType column with Spark dataframe

python

arrays

apache-spark

apache-spark-sql

pyspark