如何在PySpark中用withColumn关联其他行数据填充related_order_id

原始DataFrame #

先定义你的原始数据：

from pyspark.sql import SparkSession
from pyspark.sql.functions import when, col, first
from pyspark.sql.window import Window

spark = SparkSession.builder.appName("order_relation").getOrCreate()

data = [
    ("购买", 123, "abc", None),
    ("退货", 123, "bcd", None),
    ("购买", 234, "xyz", None),
    ("退货", 234, "zzz", None)
]

df = spark.createDataFrame(data, ["order_type", "customer_id", "order_id", "related_order_id"])
df.show()

解决方案 #

提供两种高效实现方式：

方法1：窗口函数（推荐） #

按customer_id分组，筛选每个客户对应的「购买」订单ID，再通过条件判断为退货行填充值：

# 定义窗口：按customer_id分组，按order_type排序确保购买行优先
window_spec = Window.partitionBy("customer_id").orderBy("order_type")

# 为退货行填充对应购买订单ID，购买行保持原null值
result_df = df.withColumn(
    "related_order_id",
    when(
        col("order_type") == "退货",
        first(when(col("order_type") == "购买", col("order_id"))).over(window_spec)
    ).otherwise(col("related_order_id"))
)

result_df.show()

方法2：分组聚合+关联 #

先提取每个客户的购买订单ID，再通过customer_id关联回原DataFrame：

# 提取每个客户的购买订单ID并命名
purchase_df = df.filter(col("order_type") == "购买").select("customer_id", col("order_id").alias("purchase_order_id"))

# 关联后为退货行填充值，最后清理临时列
result_df = df.join(purchase_df, on="customer_id", how="left") \
    .withColumn(
        "related_order_id",
        when(col("order_type") == "退货", col("purchase_order_id")).otherwise(col("related_order_id"))
    ) \
    .drop("purchase_order_id")

result_df.show()

最终结果 #

两种方法都会得到你期望的输出：

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