PySpark中为lead函数设置Timestamp类型默认值的实现方法

当然可以给lead函数设置Timestamp类型的默认值，你之前尝试没成功，大概率是因为没有把字面量正确转换成Timestamp类型，导致类型不匹配。

问题分析 #

lead函数返回的是Timestamp类型的列，如果直接用lit("9999-12-31 23:59")作为默认值，它的类型是字符串，和列的Timestamp类型不兼容，所以无法正确替换null。我们需要先把这个字面量转换成Timestamp类型，再传给lead的default参数。

修改后的代码 #

把你原来生成df_h的代码替换成下面的写法：

# 构造Timestamp类型的默认值
default_ts = to_timestamp(lit("9999-12-31 23:59"), "yyyy-MM-dd HH:mm")
df_h = df_ts.withColumn(
    'endts', 
    lead('startts', 1, default=default_ts).over(Window.partitionBy('id_1', 'id_2').orderBy('startts'))
)

完整验证代码 #

把整个流程整合起来，方便你直接运行验证：

import sys
from pyspark.context import SparkContext
from pyspark.sql import functions as F
from pyspark.sql.functions import *
from pyspark.sql.window import Window
from pyspark.sql.functions import lag, lead, first, last

df = spark.createDataFrame(
    [
        ("120", "aaaaa", "2021-11-19 08:00", 3.2),
        ("120", "aaaaa", "2021-11-19 10:00", 3.3),
        ("120", "aaaaa", "2021-11-20 08:00", 3.1),
        ("120", "aaaaa", "2021-11-20 15:00", 2.4),
        ("150", "bbbbb", "2021-11-19 08:45", 2.7),
        ("150", "bbbbb", "2021-11-19 09:00", 2.2),
        ("150", "bbbbb", "2021-11-19 14:00", 2.1),
        ("150", "bbbbb", "2021-11-19 21:00", 2.9),
        ("150", "bbbbb", "2021-11-19 23:00", 2.4),
        ("203", "ccccc", "2021-11-19 08:00", 2.1),
        ("203", "ccccc", "2021-11-19 09:00", 2.3),
        ("203", "ccccc", "2021-11-19 10:00", 2.4),
        ("203", "ccccc", "2021-11-19 11:00", 2.2)
    ],
    ["id_1", "id_2", "ts", "value"]
)

# 将字符串转换为Timestamp类型
df_ts = df.withColumn("startts",to_timestamp('ts', 'yyyy-MM-dd HH:mm')).drop('ts')

# 生成endts列，为最后一行指定默认Timestamp值
default_ts = to_timestamp(lit("9999-12-31 23:59"), "yyyy-MM-dd HH:mm")
df_h = df_ts.withColumn(
    'endts', 
    lead('startts', 1, default=default_ts).over(Window.partitionBy('id_1', 'id_2').orderBy('startts'))
)

# 查看最终结果
df_h.show(truncate=False)

执行结果 #

修改后每个分组的最后一行endts会被正确替换为指定的Timestamp值：

+----+-----+-----+-------------------+-------------------+
|id_1|id_2 |value|startts            |endts              |
+----+-----+-----+-------------------+-------------------+
|120 |aaaaa|3.2  |2021-11-19 08:00:00|2021-11-19 10:00:00|
|120 |aaaaa|3.3  |2021-11-19 10:00:00|2021-11-20 08:00:00|
|120 |aaaaa|3.1  |2021-11-20 08:00:00|2021-11-20 15:00:00|
|120 |aaaaa|2.4  |2021-11-20 15:00:00|9999-12-31 23:59:00|
|150 |bbbbb|2.7  |2021-11-19 08:45:00|2021-11-19 09:00:00|
|150 |bbbbb|2.2  |2021-11-19 09:00:00|2021-11-19 14:00:00|
|150 |bbbbb|2.1  |2021-11-19 14:00:00|2021-11-19 21:00:00|
|150 |bbbbb|2.9  |2021-11-19 21:00:00|2021-11-19 23:00:00|
|150 |bbbbb|2.4  |2021-11-19 23:00:00|9999-12-31 23:59:00|
|203 |ccccc|2.1  |2021-11-19 08:00:00|2021-11-19 09:00:00|
|203 |ccccc|2.3  |2021-11-19 09:00:00|2021-11-19 10:00:00|
|203 |ccccc|2.4  |2021-11-19 10:00:00|2021-11-19 11:00:00|
|203 |ccccc|2.2  |2021-11-19 11:00:00|9999-12-31 23:59:00|
+----+-----+-----+-------------------+-------------------+

关键说明 #

• lead函数的第三个参数就是default，专门用来指定当没有下一行时返回的值
• 必须保证default的值类型和startts列的类型完全一致（都是Timestamp），所以要用to_timestamp把字符串字面量转换成Timestamp类型，不能直接传字符串

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