【トリビアのDelta Lake】#9 Pysparkで、timestamp同士の計算をうまくやる【DayTimeInterval】

Spark DataFrameにTimestampデータ型のカラムがあり、それらの値の差を計算する機会がありました。

以下のように、サンプルデータで再現。startカラムとendカラムで、10秒違いますね。

sample_data = [{"start": "1970-01-01 09:00:00","end": "1970-01-01 09:00:10"}]
sdf = spark.createDataFrame(sample_data)
sdf.show()
>>>
+-------------------+-------------------+
|                end|              start|
+-------------------+-------------------+
|1970-01-01 09:00:10|1970-01-01 09:00:00|
+-------------------+-------------------+

このままだと見てくれはtimestampな値だが実際はstring型のままなので、一旦、timestampデータ型に型変換した「sdf_fix」を生成し、これを使ってみます。

from pyspark.sql.functions import *
sdf_fix = (
    sdf
    .select(
        col("start").cast("timestamp"),
        col("end").cast("timestamp")
    )
)

「DayTimeInterval」型ってなんだ？

ではend-startで計算した「diff」というカラムをwithColumnで作ってみます。

sdf_fix_calc = (
    sdf_fix
    .withColumn("diff",col("end")-col("start"))
)
sdf_fix_calc.show(truncate=False)

>>>
+-------------------+-------------------+-----------------------------------+
|start              |end                |diff                               |
+-------------------+-------------------+-----------------------------------+
|1970-01-01 09:00:00|1970-01-01 09:00:10|INTERVAL '0 00:00:10' DAY TO SECOND|
+-------------------+-------------------+-----------------------------------+

うお、なんだこれ。
ひとまずデータ型を確認します。

for i in sdf_fix_calc.schema:
    print(i)

>>>
StructField('start', TimestampType(), True)
StructField('end', TimestampType(), True)
StructField('diff', DayTimeIntervalType(0, 3), True)

どうやらtimestampデータの演算結果であるdiffは、「DayTimeInterval」という見慣れないデータ型になるようです。
まあ値も「INTERVAL '0 00:00:10' DAY TO SECOND」と、欲しかった結果である10秒ということがぱっと分からなくもないのですが…
例えばこの演算で得た値をその後計算に使用する、というケースですと不便が生じそうですよね。

一旦long型に変更

結論、一旦longなど数値系のデータ型に変更すれば、問題なかったです。
以下、成功したコードです。「diff_seconds」カラムが新設されています。

sdf_calc_diff = (
    sdf_fix
    .withColumn("diff_seconds",col("end").cast("long") - col('start').cast("long"))
)
sdf_calc_diff.show(truncate=False)
>>>
+-------------------+-------------------+------------+
|start              |end                |diff_seconds|
+-------------------+-------------------+------------+
|1970-01-01 09:00:00|1970-01-01 09:00:10|10          |
+-------------------+-------------------+------------+

ちなみにdiff_secondsのデータ型はlongです。
long同士の演算なので、結果もlongになるのは納得ですね。

for i in sdf_calc_diff.schema:
    print(i)
>>>
StructField('start', TimestampType(), True)
StructField('end', TimestampType(), True)
StructField('diff_seconds', LongType(), True)

また、longは整数系なので出来ないですが、少数も含むdoubleに型変換して演算すると、秒以下も取れるみたいです。

sdf_calc_diff = (
    sdf_fix
    .withColumn("diff_seconds",col("end").cast("double") - col('start').cast("double"))
)
sdf_calc_diff.show(truncate=False)

>>>
+-------------------+-------------------+------------+
|start              |end                |diff_seconds|
+-------------------+-------------------+------------+
|1970-01-01 09:00:00|1970-01-01 09:00:10|10.0        |
+-------------------+-------------------+------------+

僕と同じように「timestampのカラムで演算したけど、『DayTimeInterval』ってなんだよう(´°̥̥̥̥̥̥̥̥ω°̥̥̥̥̥̥̥̥｀)」と面食らった人の参考になれば幸いです〜！

参考

Pysparkのデータ型はDatabricksのドキュメントが分かりやすいです。