如何在PySpark中对日期列排序并生成列名数组列（忽略Null）

现有数据集包含person_id字段及d1、d2、d3、d4四个日期类型字段，需求为新增result_column列：忽略日期字段的Null值，将非空日期列按日期升序排序后，把对应列名以数组形式存入该列。示例数据及预期结果如下：

实现思路 #

1. 将每个日期字段与对应列名组合成结构体（日期为值，列名为标识），生成结构体数组
2. 过滤数组中日期为Null的结构体元素
3. 按结构体中的日期字段对数组进行升序排序
4. 提取排序后数组中的列名，生成最终结果数组

PySpark代码实现 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import array, struct, col, filter, sort_array, transform, lit

# 初始化SparkSession
spark = SparkSession.builder.appName("DateColumnSort").getOrCreate()

# 构造示例数据集
data = [
    ("p1", "2022-01-23", "2022-01-22", None, "2022-01-25"),
    ("p2", "2022-02-15", "2022-02-16", None, None)
]
columns = ["person_id", "d1", "d2", "d3", "d4"]
df = spark.createDataFrame(data, columns)

# 将字符串类型的日期转为date类型（若原始数据已是date类型可跳过此步）
df = df.withColumn("d1", col("d1").cast("date")) \
       .withColumn("d2", col("d2").cast("date")) \
       .withColumn("d3", col("d3").cast("date")) \
       .withColumn("d4", col("d4").cast("date"))

# 生成目标列result_column
df_result = df.withColumn(
    "date_structs",
    array(
        struct(col("d1").alias("date"), lit("d1").alias("col_name")),
        struct(col("d2").alias("date"), lit("d2").alias("col_name")),
        struct(col("d3").alias("date"), lit("d3").alias("col_name")),
        struct(col("d4").alias("date"), lit("d4").alias("col_name"))
    )
).withColumn(
    "filtered_structs",
    filter(col("date_structs"), lambda x: x["date"].isNotNull())
).withColumn(
    "sorted_structs",
    sort_array(col("filtered_structs"), asc=True, sortBy="date")
).withColumn(
    "result_column",
    transform(col("sorted_structs"), lambda x: x["col_name"])
).drop("date_structs", "filtered_structs", "sorted_structs")

# 输出结果
df_result.show(truncate=False)

代码说明 #

• struct：将日期字段和列名绑定为结构体，实现日期与列名的关联
• filter：剔除日期为Null的无效结构体
• sort_array：指定按结构体中的date字段升序排序数组
• transform：从排序后的结构体数组中提取列名字段，生成最终结果数组

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