如何在PySpark中为每个对象从两列生成转移序列字符串？

问题描述 #

我有含列x、s、d的数据集，其中s和d表示x中对象的转移关系，需要为每个x对应的对象生成完整的转移序列字符串（例如A->B->C）。我尝试了以下PySpark代码，但无法正常运行：

from pyspark.sql.functions import udf
from pyspark.sql.functions import array_distinct
from pyspark.sql.types import ArrayType, StringType

create_transition = udf(lambda x: "->".join([i[0] for i in groupby(x)]))

df= df\
      .withColumn('list', F.concat(df['s'], F.lit(','), df['d']))\
      .groupBy('x').agg(F.collect_list('list').alias('list2'))\
      .withColumn("list3", create_transition("list2"))

问题分析 #

原代码存在3个核心问题：

• 未导入pyspark.sql.functions别名F，导致F.concat调用报错
• 未定义groupby方法，且逻辑错误——转移序列需要按s与d的关联关系构建链式路径，而非对字符串列表分组
• UDF逻辑无法解析s和d的转移关联，直接拼接字符串列表无法生成正确的转移链

解决方案 #

方案1：纯PySpark实现（无第三方依赖） #

通过迭代关联逐步构建每个x对应的转移链，适合线性转移场景：

from pyspark.sql import functions as F
from pyspark.sql.window import Window

# 1. 定位每个x对应的转移起点（无前置输入的节点）
start_nodes = df.groupBy("x", "s").agg(F.count("d").alias("cnt"))\
    .join(df.groupBy("x", "d").agg(F.count("s").alias("cnt")), 
          on=["x", "s"], how="left_outer")\
    .filter(F.col("cnt_right").isNull())\
    .select("x", "s").withColumnRenamed("s", "current_node")

# 2. 迭代拼接转移链
def build_chain(df, start_df, max_iter=10):
    chain_df = start_df.withColumn("chain", F.col("current_node"))
    for _ in range(max_iter):
        chain_df = chain_df.join(df, 
                                 (chain_df["x"] == df["x"]) & (chain_df["current_node"] == df["s"]),
                                 how="left_outer")\
                           .withColumn("new_node", F.coalesce(df["d"], F.col("current_node")))\
                           .withColumn("new_chain", F.when(df["d"].isNotNull(), 
                                                           F.concat(F.col("chain"), F.lit("->"), df["d"]))
                                                   .otherwise(F.col("chain")))\
                           .drop("s", "d")\
                           .withColumnRenamed("new_node", "current_node")\
                           .withColumnRenamed("new_chain", "chain")
    # 保留每个x对应的最长完整转移链
    window = Window.partitionBy("x").orderBy(F.length("chain").desc())
    return chain_df.withColumn("rank", F.row_number().over(window))\
                   .filter(F.col("rank") == 1)\
                   .drop("current_node", "rank")

# 生成最终结果
result_df = build_chain(df, start_nodes)
result_df.show()

方案2：使用GraphFrames（适合复杂转移场景） #

若环境允许安装第三方库，利用图结构路径查找功能更高效，支持分支、循环类复杂转移关系：

from graphframes import GraphFrame

# 1. 构建图结构
vertices = df.select(F.col("s").alias("id")).union(df.select(F.col("d").alias("id"))).distinct()
edges = df.select(F.col("s").alias("src"), F.col("d").alias("dst"), "x")

# 2. 按x分组生成最长转移路径
def get_longest_path(graph, x_val):
    sub_graph = graph.filterEdges(F.col("x") == x_val)
    # 定位转移起点
    start_node = sub_graph.edges.groupBy("src").count()\
        .join(sub_graph.edges.groupBy("dst").count(), on="src", how="left_outer")\
        .filter(F.col("count_right").isNull())\
        .select("src").first()[0]
    # 递归拼接完整路径
    path = [start_node]
    current = start_node
    while True:
        next_node = sub_graph.edges.filter(F.col("src") == current).select("dst").first()
        if not next_node:
            break
        path.append(next_node[0])
        current = next_node[0]
    return "->".join(path)

# 生成结果数据集
x_list = df.select("x").distinct().rdd.flatMap(lambda x: x).collect()
result_rows = [(x_val, get_longest_path(GraphFrame(vertices, edges), x_val)) for x_val in x_list]
result_df = spark.createDataFrame(result_rows, ["x", "transition_sequence"])
result_df.show()

方案说明 #

• 方案1无需额外依赖，实现简单，适合无分支的线性转移场景
• 方案2需提前安装GraphFrames（pip install graphframes），处理复杂转移关系效率更高

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