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

【问题标题】：Map values in ArrayType column with Spark dataframe使用 Spark 数据框映射 ArrayType 列中的值
【发布时间】：2023-02-02 13:24:03
【问题描述】：

我有一个带有 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 亿条记录）。

【问题讨论】：

a | [b, c]怎么会变成a | [[a,d],[a]]？
在我的初始数据框中，我有 id 和邻居：b | [a, d] 和 c | [a]，所以我将这些值映射到数组 [b, c] 上id==id 内部邻居@pltc

标签： python arrays apache-spark pyspark apache-spark-sql

【解决方案1】：

您需要explode“邻居”列，然后只需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]]|
# +---+-------------+

【讨论】：