在 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 在技术上有效）。然而，闰秒是出了名的难，因为没有确切的公式（但是 - 谁知道呢，也许有一天我们可以模拟地质和气候事件？）。

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

Add days to element inside array in PySpark Dataframe

python

pyspark

pyspark-sql

pyspark-dataframes