向量化tz_convert的熊猫应用函数

时间:2020-04-20 19:32:10

标签: python pandas vectorization apply

我有一个数据框,其中hour列包含UTC中的日期时间数据。我有一个time_zone列,其中包含每个观察的时区,并且我将其用于将hour转换为本地时间并将其保存在名为local_hour的新列中。为此,我使用以下代码:

import pandas as pd

# Sample dataframe
import pandas as pd
df = pd.DataFrame({
    'hour': ['2019-01-01 05:00:00', '2019-01-01 07:00:00', '2019-01-01 08:00:00'],
    'time_zone': ['US/Eastern', 'US/Central', 'US/Mountain']
})

# Ensure hour is in datetime format and localized to UTC
df['hour'] = pd.to_datetime(df['hour']).dt.tz_localize('UTC')

# Add local_hour column with hour in local time 
df['local_hour'] = df.apply(lambda row: row['hour'].tz_convert(row['time_zone']), axis=1)

df
    hour                        time_zone   local_hour
0   2019-01-01 05:00:00+00:00   US/Eastern  2019-01-01 00:00:00-05:00
1   2019-01-01 07:00:00+00:00   US/Central  2019-01-01 01:00:00-06:00
2   2019-01-01 08:00:00+00:00   US/Mountain 2019-01-01 01:00:00-07:00

该代码有效。但是使用apply的运行速度很慢,因为实际上我有一个很大的数据框。有没有一种方法可以对此进行矢量化或以其他方式加快速度?

注意:我曾尝试使用swifter软件包,但就我而言,它不能加快速度。

1 个答案:

答案 0 :(得分:1)

从没有无限个time_zone的假设出发,也许您可​​以为每个组执行tz_convert,例如:

df['local_hour'] = df.groupby('time_zone')['hour'].apply(lambda x: x.dt.tz_convert(x.name))
print (df)

                       hour    time_zone                 local_hour
0 2019-01-01 05:00:00+00:00   US/Eastern  2019-01-01 00:00:00-05:00
1 2019-01-01 07:00:00+00:00   US/Central  2019-01-01 01:00:00-06:00
2 2019-01-01 08:00:00+00:00  US/Mountain  2019-01-01 01:00:00-07:00

在样本上,它可能会比您做的要慢,但是在更大的数据和组上,它应该会更快

为了进行速度比较,使用您提供的3行中的df,可以得出:

%timeit df.apply(lambda row: row['hour'].tz_convert(row['time_zone']), axis=1)
# 1.6 ms ± 102 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit df.groupby('time_zone')['hour'].apply(lambda x: x.dt.tz_convert(x.name))
# 2.58 ms ± 126 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

所以apply更快,但是如果您创建的数据帧大1000倍,但只有3个time_zones,那么groupby的速度大约快20倍:

df = pd.concat([df]*1000, ignore_index=True)

%timeit df.apply(lambda row: row['hour'].tz_convert(row['time_zone']), axis=1)
# 585 ms ± 42.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit df.groupby('time_zone')['hour'].apply(lambda x: x.dt.tz_convert(x.name))
# 27.5 ms ± 2.15 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)