PySpark DataFrame基于userid和transactiontime填充location空值

首先看你的原始DataFrame结构：

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

spark = SparkSession.builder.getOrCreate()
sc = spark.sparkContext

df = sc.parallelize([(123, 2345,25,""), (123, 2345,29,"NY"), (123,5422,67,"NY"),(123,9422,67,"NY"),(123,3581,98,"NY"),(231, 4322,77,""),(231,4322,99,"Paris"),(231,8342,45,"Paris")]).toDF(["userid", "transactiontime","zip","location"])

你之前的Window函数方案失败有两个核心原因：

• last('location')默认不忽略空值，当窗口内最后一行是空值时会返回空；
• 默认窗口范围是分区起始到当前行，加上按zip排序后，空值行可能排在前面，last无法取到后续的非空值。

下面提供两种可行的解决方案：

方法一：优化Window函数 #

调整窗口范围覆盖整个分区，同时设置ignoreNulls=True，确保能取到分组内任意位置的非空值：

# 定义窗口：按userid和transactiontime分区，范围覆盖整个分组
w = Window.partitionBy('userid', 'transactiontime').rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)

# 用last函数忽略空值填充，first函数效果一致（同一分组内非空location唯一）
df_filled = df.withColumn(
    'location',
    F.last('location', ignoreNulls=True).over(w)
)

df_filled.show()

方法二：分组聚合后关联 #

先提取每个分组的非空location，再和原表关联填充：

# 分组获取每个userid+transactiontime对应的非空location
location_ref = df.filter(F.col('location') != '') \
                 .groupBy('userid', 'transactiontime') \
                 .agg(F.first('location').alias('filled_location'))

# 关联原表，用coalesce优先保留原非空值，空值则用关联后的填充值
df_filled = df.join(location_ref, on=['userid', 'transactiontime'], how='left') \
              .withColumn('location', F.coalesce(F.col('location'), F.col('filled_location'))) \
              .drop('filled_location')

df_filled.show()

两种方法都会得到预期输出：

+------+---------------+---+--------+
|userid|transactiontime|zip|location|
+------+---------------+---+--------+
|   123|           2345| 25|      NY|
|   123|           2345| 29|      NY|
|   123|           5422| 67|      NY|
|   123|           9422| 67|      NY|
|   123|           3581| 98|      NY|
|   231|           4322| 77|    Paris|
|   231|           4322| 99|    Paris|
|   231|           8342| 45|    Paris|
+------+---------------+---+--------+

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