pandas, как эффективно разбить большой фрейм данных на два набора с сгруппированным условием на datetime

Question

У меня большой фрейм данных (~ 40 миллионов строк), и я хочу разделить его на две части.Столбец 'group' указывает, к какой группе относится образец, а в столбце 'date' - какая дата выборки.В следующем тестовом примере может быть несколько одинаковых выборок, но в исходном наборе это не так.Далее в исходных данных есть три столбца.Теперь, если самая поздняя дата одной группы старше или равна '2017-01-30', я хочу установить для всей группы значение s1, в противном случае - s2.

Приведенный ниже код делает то, что должен, ноэто очень медленно.У вас есть идея, как я могу ускорить процесс?Знаете ли вы, почему этот подход такой медленный?Линия df['split'] = grouped['date'].transform(lambda x: x.max() < date_) является узким местом.Спасибо за вашу помощь.

from random import randint
import numpy as np
import time
import pandas as pd

length = int(1e5)

bimonthly_days = np.arange(0, 30)
base_date = np.datetime64('2017-01-01')
random_date = base_date + np.random.choice(bimonthly_days)

groups = np.random.randint(1, int(2e4), length)
dates = np.array([base_date + np.random.choice(bimonthly_days) for _ in range(length)], dtype='datetime64[ns]')

df = pd.DataFrame({'group': groups, 'date': dates})

grouped = df.groupby('group')
date_ = np.datetime64('2017-01-30')

start_time = time.process_time()
df['split'] = grouped['date'].transform(lambda x: x.max() < date_)
dif = time.process_time() - start_time
print(f" elapsed time: {dif}")

s1 = df[df['split'] == 1].drop(columns=['split'])
s2 = df[df['split'] == 0].drop(columns=['split'])

Вывод:

elapsed time: 12.923806

Answer 1

Попробуйте это:

start_time = time.process_time() df.loc[:,'split'] = df.groupby('group').date.transform('max') dif = time.process_time() - start_time print(f" elapsed time: {dif}")

s1 = df[df.split < date_] s2 = df[df.split >= date_]

 elapsed time: 0.01961299999999966

С вашим методом мое время:

elapsed time: 6.9499130000000005

Answer 2

Быстрее создать Series с transform, а затем сравнить только один раз, как сравнивать каждую группу в отдельности.Также вы можете создать новую маску вместо столбца для сравнения, а для инвертированной маски используйте ~:

mask = grouped['date'].transform('max') < date_

s1 = df[mask]
s2 = df[~mask]

Производительность :

np.random.seed(10)
length = int(1e5)

bimonthly_days = np.arange(0, 30)
base_date = np.datetime64('2017-01-01')
random_date = base_date + np.random.choice(bimonthly_days)

groups = np.random.randint(1, int(2e4), length)
dates = np.array([base_date + np.random.choice(bimonthly_days) for _ in range(length)], dtype='datetime64[ns]')

df = pd.DataFrame({'group': groups, 'date': dates})


In [219]: %timeit df.groupby('group')['date'].transform('max') < np.datetime64('2017-01-30')
9.81 ms ± 645 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [220]: %timeit df.groupby('group')['date'].transform(lambda x: x.max() < np.datetime64('2017-01-30'))
9.05 s ± 159 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)