PySpark如何从含节点数组列的DataFrame生成相邻节点配对表

基于你现有代码的实现方案 #

你编写的UDF返回的是字典数组，只需要先明确指定UDF的返回schema，再配合explode算子拆解数组即可得到目标结构：

步骤1：定义UDF返回schema

from pyspark.sql.types import StructType, StructField, StringType, LongType, MapType, ArrayType
from pyspark.sql.functions import udf, explode, col

# 定义返回数组中每个弧的结构
arc_schema = ArrayType(
    StructType([
        StructField("id", StringType(), nullable=False),
        StructField("from", LongType(), nullable=False),
        StructField("to", LongType(), nullable=False),
        StructField("tags", MapType(StringType(), StringType()), nullable=True)
    ])
)

def split_ways_into_arcs(row):
    arcs = []
    for node in range(len(row['nodes']) - 1):
        arc = dict()
        arc['id'] = str(row['nodes'][node]) + "_" + str(row['nodes'][node + 1])
        arc['from'] = row['nodes'][node]
        arc['to'] = row['nodes'][node + 1]
        arc['tags'] = row['tags']
        arcs.append(arc)
    return arcs

# 声明UDF时指定返回schema
split = udf(lambda row: split_ways_into_arcs(row.asDict()), returnType=arc_schema)

步骤2：生成目标DataFrame

result_df = osmDF\
    .withColumn("arcs", split(col("*")))  # 传入整行数据调用UDF，得到弧数组
    .select("*", explode("arcs").alias("arc"))  # 把弧数组拆分为多行
    .select(  # 提取结构体字段得到最终结构
        col("arc.id").alias("id"),
        col("arc.from").alias("from"),
        col("arc.to").alias("to"),
        col("arc.tags").alias("tags")
    )

步骤3：保存为parquet

result_df.write.parquet("目标保存路径", mode="overwrite")

更优的原生算子实现方案（无UDF，性能更高） #

完全使用Spark内置算子实现，避免自定义UDF的序列化开销，适合大规模数据场景：

from pyspark.sql.functions import posexplode, concat_ws, size, array_slice

result_df = osmDF\
    .select(
        "*",
        # 遍历节点的前n-1个元素，得到下标和起始节点
        posexplode(array_slice(col("nodes"), 1, size(col("nodes")) - 1)).alias("idx", "from")
    )\
    .withColumn("to", col("nodes")[col("idx") + 1])\
    .withColumn("id", concat_ws("_", col("from"), col("to")))\
    .select("id", "from", "to", "tags")

内容的提问来源于stack exchange，提问作者Laurent