如何優化 df.apply 函數

Question

我有以下格式的兩個數據框。

數據幀 A：

DateTime            | A |
-------------------------
2020-01-01 06:34:12 | 1 |
2020-01-01 06:36:24 | 2 |
2020-01-01 06:36:28 | 3 |
...

數據幀 B：

StartDateTime       | EndDateTime         | Value |
---------------------------------------------------
2020-01-01 06:30:00 | 2020-01-01 06:35:00 | 1.5   |
2020-01-01 06:35:00 | 2020-01-01 06:40:00 | 1.2   |
...

最終，我想通過從 DataFrame A 中獲取 DateTime 並找到日期時間在 StartDateTime 和 EndDateTime 之間的行，按如下方式組合兩個數據幀：

DateTime            | A | Value |
---------------------------------
2020-01-01 06:34:12 | 1 | 1.5   |
2020-01-01 06:36:24 | 2 | 1.2   |
2020-01-01 06:36:28 | 3 | 1.2   |
...

我正在使用以下內容，但速度非常慢：

df_a['Value'] = df_a['DateTime'].apply(lambda x: df_b.loc[(df_b['StartDateTime'] <= x) & (df_b['EndDateTime'] > x)]['Value'].iloc[0])

我應該如何重寫這個，因為我有 1MM+ 行的數據幀，而且它目前非常慢。

Answer 1

情況 1：bin 的大小可變

在這種情況下，我能想到的最好的方法是使用pd.cut ：

mapper = pd.Series(df_b['Value'])
mapper.index = df_b['StartDateTime']

cutoffs = df_b['StartDateTime'].copy()
cutoffs[cutoffs.index.max() + 1] = df_b['EndDateTime'].max()

bins = pd.cut(df_a['DateTime'], bins=cutoffs)
df_a['Value'] = mapper.loc[pd.IntervalIndex(bins).left].values

您創建一個系列以將開始時間映射到值。 然后創建另一個系列，表示來自 DataFrame A 的時間將被分箱到其中的截止點（請注意，您需要手動添加最后一個結束時間）。 然后你斌時代與那些截止pd.cut ，並使用left二進位值loc映射系列。

情況 2：垃圾箱大小相同

看起來 OP 的垃圾箱是 5 分鍾的大塊。 如果這是對的，您可以利用pd.Series.dt.floor()將時間從 DataFrame A 快速轉換為可以索引 DataFrame B 的時間：

mapper = pd.Series(df_b['Value'])
mapper.index = df_b['StartDateTime']
df_a['Value'] = mapper.loc[df_a['DateTime'].dt.floor('5T')].values

定時：

這是我使用的示例數據：

import numpy as np
import pandas as pd

size = 100 # tweak this to see each option at scale

dr1 = pd.date_range('01-01-2020 06:00:00', freq='5T', periods=size)
dr2 = pd.date_range('01-01-2020 06:05:00', freq='5T', periods=size)
drA = pd.to_datetime({'year':dr1.year, 'month':dr1.month,
                      'day':dr1.day, 'hour':dr1.hour,
                      'minute':np.random.randint(1,60,len(dr1)),
                      'second':np.random.randint(1,60,len(dr1))}).sort_values()
drA = drA[drA < dr2.max()]

df_a = pd.DataFrame({'DateTime':drA, 'A':range(len(drA))})
df_b = pd.DataFrame({'StartDateTime':dr1, 'EndDateTime':dr2, 'Value':np.random.rand(len(dr2))})

使用%%timeit和size=100 ：

• apply ：每個循環 61 ms ± 851 µs（平均值 ± 標准偏差，7 次運行，每次 10 次循環）
• pd.cut ：每個循環 8.98 ms ± 107 µs（平均值 ± 標准pd.cut ，7 次運行，每次 100 次循環）
• dt.floor ：每個循環 865 µs ± 17.8 µs（平均值 ± 標准dt.floor ，7 次運行，每次 1000 次循環）
• 使用np.where *添加 @Rik Kraan 的答案：每個循環 1.85 ms ± 7.8 µs（平均值 ± 標准np.where ，7 次運行，每次 1000 次循環）

*這個答案比我的pd.cut ，但是當將size增加到1000000時，我也得到了一個MemoryError: Unable to allocate 931. GiB for an array with shape (999999, 1000000) and data type bool pd.cut MemoryError: Unable to allocate 931. GiB for an array with shape (999999, 1000000) and data type bool 。

所以發言是顯著比原來的方法快。 但是，如果您的垃圾箱不是平均分配的，那將是不正確的。 您可以使用df_b['StartDateTime'].dt.minute.unique()或df_b['StartDateTime'].dt.time.unique() 。 如果您能找到合適的值，甚至可以反復使用多個樓層值。

pd.cut版本仍然是一個顯着的改進； 也許還有一些我沒有看到的其他優化？

Answer 2

讓我們首先創建兩個數組，返回滿足條件的兩個dfs A和B的索引（ A['DateTime']介於B['StartDateTime'] & B['EndDateTime']

i, j = np.where(
(A['DateTime'].values[:, None] >= B['StartDateTime'].values) & 
(A['DateTime'].values[:, None] <= B['EndDateTime'].values)
)

從與這些索引對應的數據幀A和B選擇行並創建一個新的數據幀

pd.DataFrame(
    np.column_stack([A.values[i], B.values[j]]),
    columns=A.columns.append(B.columns)
)