python - 在 PySpark Dataframe 中向数组内的元素添加天数

Question

我有一个包含三列的 PySpark 数据框。前两列以数组作为元素，最后一列给出最后一列的数组长度。以下是 PySpark 数据框:

+---------------------+---------------------+-----+
|                   c1|                   c2|lenc2|
+---------------------+---------------------+-----+
|[2017-02-14 00:00:00]|[2017-02-24 00:00:00]|    1|
|[2017-01-16 00:00:00]|                   []|    0|
+---------------------+---------------------+-----+

数组包含时间戳数据类型。 lenc2 列表示 c1 列中数组的长度。对于 lenc2==0 的所有行，列 c1 只有一个(时间戳)元素。

对于 lenc2==0 的所有行，我想从列 c1 中的数组中获取时间戳，添加 5 天并将其放入在 c2 行的数组内。我怎样才能做到这一点？

这是预期输出的示例:

+---------------------+---------------------+-----+
|                   c1|                   c2|lenc2|
+---------------------+---------------------+-----+
|[2017-02-14 00:00:00]|[2017-02-24 00:00:00]|    1|
|[2017-01-16 00:00:00]|[2017-01-21 00:00:00]|    0|
+---------------------+---------------------+-----+

以下是我迄今为止尝试过的:

df2 = df1.withColumn(
    "c2",
    F.when(F.col("lenc2") == 0, F.array_union(F.col("c1"), F.col("c2"))).otherwise(
        F.col("c2")
    ),
)

Answer 1

您的 when(...).otherwise(...) 已经正确。

鉴于您似乎对亚秒精度不感兴趣，您可以将时间戳转换为自 Unix 纪元以来的秒数，并添加 5 天的秒数，然后转换回时间戳:

from datetime import datetime

from pyspark.sql.functions import *

one_sec_before_leap_time = datetime(2016, 12, 31, 23, 59, 59)
seconds_in_a_day = 24 * 3600

df = spark.createDataFrame([
    ([one_sec_before_leap_time], [datetime.now()], 1),
    ([one_sec_before_leap_time], [], 0),
],
    schema=("c1", "c2", "lenc2"))


def add_seconds_to_timestamp(ts_col, seconds_col):
    return to_timestamp(unix_timestamp(ts_col) + seconds_col)


df2 = df.withColumn("c2",
                    when(col("lenc2") == 0,
                         array(
                             add_seconds_to_timestamp(
                                 col("c1").getItem(0),
                                 lit(5 * seconds_in_a_day))))
                    .otherwise(col("c2")))
df2.show(truncate=False)
# +---------------------+----------------------------+-----+                      
# |c1                   |c2                          |lenc2|
# +---------------------+----------------------------+-----+
# |[2016-12-31 23:59:59]|[2019-12-07 16:58:32.864176]|1    |
# |[2016-12-31 23:59:59]|[2017-01-05 23:59:59]       |0    |
# +---------------------+----------------------------+-----+

请注意，当您必须考虑夏令时时，这很可能会给您带来奇怪的结果。最好用 UTC 表示所有内容，并且仅在输入和输出处进行从 UTC 时间戳到以本地时区表示的时间的适当转换。基本上类似于 Unicode 三明治。

此外，这没有考虑leap seconds ，如上所示(2016 年还有一秒钟，使得 2016-12-31T12:59:60Z 在技术上有效)。然而，闰秒是出了名的困难，因为它没有确切的公式(但是 - 谁知道，也许有一天我们可以模拟地质和气候事件？)。