加快新特征(时间和对象的行数循环)的计算速度

Question

我有一个包含汽车及其价格的数据框，其中每行包含汽车的售价和销售时间。我想创建一个新特性，每行显示该车在过去2小时内的销量乘以0.2。

res = pd.DataFrame()
for car in pd.unique(df['car_id']):

car_rows = df[df['car_id'] == car]

sum_cars = []
for row in car_rows.itertuples():
rows = car_rows[(car_rows['time'] < row.time) & (car_rows['time'] > row.time - pd.Timedelta(hours=2))]
if rows.shape[0] == 0:
sum_cars.append(0)
continue
sum_cars.append(np.sum(rows.apply(lambda x: 0.2 * x.price, axis=1)))
car_rows['sum_cars'] = sum_cars
res = pd.concat([res, car_rows])

这段代码非常慢(30分钟500k行)。有什么加快速度的办法吗?

示例(计算特征sum_cars):

car_id  time               price    sum_cars
1   2019-01-02 16:55:13  10    0.0
1   2019-01-02 16:59:37  20    2.0
1   2019-01-02 17:35:18  30    6.0
2   2019-01-03 17:45:13  40    0.0
2   2019-01-03 18:35:13  50    8.0
2   2019-01-03 20:55:13  60    0.0
2   2019-01-03 21:54:13  70    12.0

Answer 1

在pandas中有一个非常有用的结构叫做" Rolling GroupBy"(pandas.core.window.rolling.RollingGroupby，由df.groupby(...).rolling(...)得到)

对于一些包含500K行的合成数据，我看到执行时间大约为135ms(加速超过13000倍)。(注意:这是使用pandas=1.2.1;一个仍然使用pandas=1.1.4的朋友观察到的速度比我在这里报告的慢35倍——这是一个很好的提醒，要定期更新)。

要计算的数量为:

>>> z = df.groupby('car_id').rolling(
...     '2H', closed='left', min_periods=0)['price'].sum() * 0.2
>>> z
car_id  time               
1       2019-01-02 16:55:13     0.0
2019-01-02 16:59:37     2.0
2019-01-02 17:35:18     6.0
2       2019-01-03 17:45:13     0.0
2019-01-03 18:35:13     8.0
2019-01-03 20:55:13     0.0
2019-01-03 21:54:13    12.0

现在，技巧是重新排列df，使其索引为(car_id, time)(我们也对其进行排序)。然后我们可以将计算出的数量z赋值给新列sum_cars:

out = df.set_index(
'car_id', append=True).swaplevel().sort_index().assign(sum_cars=z)

:

>>> out
price  sum_cars
car_id time                                
1      2019-01-02 16:55:13     10       0.0
2019-01-02 16:59:37     20       2.0
2019-01-02 17:35:18     30       6.0
2      2019-01-03 17:45:13     40       0.0
2019-01-03 18:35:13     50       8.0
2019-01-03 20:55:13     60       0.0
2019-01-03 21:54:13     70      12.0

def gen_data(n=10):
df = pd.DataFrame({
'car_id': np.random.randint(0, 3, n),
'price': np.random.normal(10, 5, n).round(),
'time': pd.Series(
pd.Timestamp('2019-01-02') + np.random.uniform(0, 5, n)
* pd.Timedelta('1H')).dt.floor('s'),
}).set_index('time').sort_index()
return df
def func(df):
z = df.groupby('car_id').rolling(
'2H', closed='left', min_periods=0)['price'].sum() * 0.2
return (
df.set_index('car_id', append=True).swaplevel().sort_index()
.assign(sum_cars=z)
)
df = gen_data(500_000)
%timeit func(df)
# 135 ms ± 173 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)