Rounding datetime to nearest midnight using Spark SQL

Viewed 933

I have a Spark dataframe with datetime column in format yyyy-MM-dd HH:mm:SS.

My dataframe looks like below:

datetime_column output_column
2021-01-29 21:30:00 2021-01-30 00:00:00
2021-01-30 05:30:00 2021-01-30 00:00:00
2021-01-01 23:25:00 2021-01-02 00:00:00
2021-01-02 08:59:59 2021-01-02 00:00:00

How would I be able to get the output_column column from datetime_column column using Spark SQL?

NOTE: My data (datetime_column) will always be present between 9 PM to 9 AM so finding the nearest midnight is straightforward.

I am using PySpark 2.4.6

3 Answers

Here's a trick that performs that "rounding" by computing the nearest day in second values (i.e, round((timestamp // 86400) * 86400). It's the same thing I would do to round a number to the nearest 1000: round(number//1000)*1000)

(df.select(f.from_unixtime(f.round(
    f.unix_timestamp('datetime_column') / 
    f.lit(86400)) * f.lit(86400)).alias('output_column') )).show()

Outputs:

+-------------------+
|      output_column|
+-------------------+
|2021-01-30 02:00:00|
|2021-01-30 02:00:00|
|2021-01-02 02:00:00|
|2021-01-02 02:00:00|
+-------------------+

Note that the result above is shown in my time zone, whose offset is +02:00 (from_unixtime converts into a time stamp of the current time zone). If you need to have your exact result, be sure to set UTC as system time zone on your Spark config before running the code (this is mentioned in the function's docs):

spark.conf.set("spark.sql.session.timeZone", "UTC")
(df.select(f.from_unixtime(f.round(
    f.unix_timestamp('datetime_column') / 
    f.lit(86400)) * f.lit(86400)).alias('output_column'))).show()
spark.conf.unset("spark.sql.session.timeZone") # unset it

This results in your exact expected output

+-------------------+
|      output_column|
+-------------------+
|2021-01-30 00:00:00|
|2021-01-30 00:00:00|
|2021-01-02 00:00:00|
|2021-01-02 00:00:00|
+-------------------+

A more intuitive way of doing this - you can use a case when to check the hour, and add 1 to the date if the hour is later than 9 PM. Otherwise keep the date and remove the time.

import pyspark.sql.functions as F

df2 = df.withColumn(
    'output_column',
    F.when(
        F.hour('datetime_column') >= 21,
        F.date_add(F.date_trunc('day', 'datetime_column'), 1)
    ).otherwise(F.date_trunc('day', 'datetime_column'))
)

df2.show()
+-------------------+-------------------+
|    datetime_column|      output_column|
+-------------------+-------------------+
|2021-01-29 21:30:00|2021-01-30 00:00:00|
|2021-01-30 05:30:00|2021-01-30 00:00:00|
|2021-01-01 23:25:00|2021-01-02 00:00:00|
|2021-01-02 08:59:59|2021-01-02 00:00:00|
+-------------------+-------------------+

If you prefer Spark SQL:

df.createOrReplaceTempView('df')

df2 = spark.sql("""
    select *,
        case when hour(datetime_column) >= 21
        then date_add(date_trunc('day', datetime_column), 1)
        else date_trunc('day', datetime_column)
        end as output_column 
    from df
""")

df2.show()
+-------------------+-------------------+
|    datetime_column|      output_column|
+-------------------+-------------------+
|2021-01-29 21:30:00|2021-01-30 00:00:00|
|2021-01-30 05:30:00|2021-01-30 00:00:00|
|2021-01-01 23:25:00|2021-01-02 00:00:00|
|2021-01-02 08:59:59|2021-01-02 00:00:00|
+-------------------+-------------------+

I think the simplest way to round a timestamp to the nearest date is to add half a day (12 hours) then truncate the resulting timestamp to DD unit (day) using date_trunc. This will work independently from the time when your data will be present:

from pyspark.sql import functions as F

df1 = df.withColumn(
    "output_column",
    F.date_trunc("DD", F.col("datetime_column") + F.expr("INTERVAL 12 HOURS"))
)

df1.show(truncate=False)

#+-------------------+-------------------+
#|datetime_column    |output_column      |
#+-------------------+-------------------+
#|2021-01-29 21:30:00|2021-01-30 00:00:00|
#|2021-01-30 05:30:00|2021-01-30 00:00:00|
#|2021-01-01 23:25:00|2021-01-02 00:00:00|
#|2021-01-02 08:59:59|2021-01-02 00:00:00|
#+-------------------+-------------------+

Using plain Spark-SQL:

SELECT DATE_TRUNC('DD', datetime_column  + INTERVAL 12 HOURS)
Related