Pandas: индекс не сохраняется при выполнении групповой прокрутки даты и времени

Question

У меня есть датафрейм, в котором некоторые даты совпадают. В качестве иллюстрации проблемы я создал образец df с одинаковыми датами.

df = pd.DataFrame({"column1": range(6), 
                   "column2": range(6), 
                   'group': 3*['A','B'], 
                   'date':pd.date_range("20190101", periods=6)})
df.loc[:,'date']=df.loc[0,'date']

df
# Output of DF
   column1  column2 group       date
0        0        0     A 2019-01-01
1        1        1     B 2019-01-01
2        2        2     A 2019-01-01
3        3        3     B 2019-01-01
4        4        4     A 2019-01-01
5        5        5     B 2019-01-01

Проблема возникает при выполнении операции группирования путем прокрутки столбца datetime: индекс не сохраняется. Это проблема, когда даты совпадают, так как нет возможности выполнить слияние с исходным фреймом данных (это то, к чему я стремлюсь).

df.groupby('group').rolling('2D',on='date')['column1'].sum()

# Output of Groupby Rolling
group  date      
A      2019-01-01    0.0
       2019-01-01    2.0
       2019-01-01    6.0
B      2019-01-01    1.0
       2019-01-01    4.0
       2019-01-01    9.0
Name: column1, dtype: float64

У меня есть альтернативное решение, которое работает, но значительно медленнее.

df.groupby('group').apply(lambda x: x.rolling('2D',on='date')['column1'].sum())

# Output of Groupby Apply Rolling 
group   
A      0    0.0
       2    2.0
       4    6.0
B      1    1.0
       3    4.0
       5    9.0
Name: column1, dtype: float64

Надеемся на что-то более эффективное, чем указано выше.

Answer 1

Для тех, кто заинтересован, я создал более сложный образец df для проверки эффективности каждого из предложенных выше решений.

Мой оригинальный подход (здесь самый медленный, но эффективный, если групп очень мало):

%%timeit
df = pd.DataFrame({"column1": range(600), 
                   "column2": range(600), 
                   "column3": range(600),
                   "column4": range(600),
                   "column5": range(600),
                   "column6": range(600),
                   "column7": range(600),
                   "column8": range(600),
                   'group': 5*['l'+str(i) for i in range(120)], 
                   'date':pd.date_range("20190101", periods=600)})

### Set the date the same
df.loc[:,'date']=df.loc[0,'date']

cols = ['column1','column2','column3','column4','column5','column6','column7','column8']
newcols = ['col1','col2','col3','col4','col5','col6','col7','col8']
if newcols[0] not in df.columns:
    df = df.reindex(columns=df.columns.tolist()+newcols)

df[newcols]=df.groupby('group').apply(lambda x: x.rolling('2D',on='date')[cols].sum()
                                     ).sort_index(level=1).drop('date',axis=1).values

# timeit output
345 ms ± 28 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Решение Дэвида Эриксона. Это эффективно, если в каждой группе много групп с небольшим количеством наблюдений.

%%timeit
df = pd.DataFrame({"column1": range(600), 
                   "column2": range(600), 
                   "column3": range(600),
                   "column4": range(600),
                   "column5": range(600),
                   "column6": range(600),
                   "column7": range(600),
                   "column8": range(600),
                   'group': 5*['l'+str(i) for i in range(120)], 
                   'date':pd.date_range("20190101", periods=600)})

### Set the date the same
df.loc[:,'date']=df.loc[0,'date']

cols = ['column1','column2','column3','column4','column5','column6','column7','column8']
newcols = ['col1','col2','col3','col4','col5','col6','col7','col8']
if newcols[0] not in df.columns:
    df = df.reindex(columns=df.columns.tolist()+newcols)

my_dict = {}
my_dict["index"] = "max"
my_dict.update(dict.fromkeys(cols, "sum"))
df[newcols]=df.reset_index().groupby('group').rolling('2D',
on='date').agg(my_dict).sort_values('index').drop('index',axis=1).values

# timeit output
110 ms ± 11.5 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Решение, предложенное RichieV, которое здесь самое быстрое:

%%timeit
df = pd.DataFrame({"column1": range(600), 
                   "column2": range(600), 
                   "column3": range(600),
                   "column4": range(600),
                   "column5": range(600),
                   "column6": range(600),
                   "column7": range(600),
                   "column8": range(600),
                   'group': 5*['l'+str(i) for i in range(120)], 
                   'date':pd.date_range("20190101", periods=600)})

### Set the date the same
df.loc[:,'date']=df.loc[0,'date']

cols = ['column1','column2','column3','column4','column5','column6','column7','column8']
newcols = ['col1','col2','col3','col4','col5','col6','col7','col8']
if newcols[0] not in df.columns:
    df = df.reindex(columns=df.columns.tolist()+newcols)
    
df=df.sort_values(['group','date'],kind='mergesort').reset_index(drop=True)
df[newcols]=df.groupby('group').rolling('2D',on='date')[cols].sum().values
df=df.sort_values('column1',kind='mergesort').reset_index(drop=True)

# timeit output
40 ms ± 6.41 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Answer 2

Вы можете использовать .reset_index(), а затем включить этот столбец index как результат остальных с .groupby и .agg. Я предполагаю, что это будет намного быстрее, чем лямбда x.

df = pd.DataFrame({"column1": range(6), 
               "column2": range(6), 
               'group': 3*['A','B'], 
               'date':pd.date_range("20190101", periods=6)})
df = df.reset_index().groupby('group').rolling('5D',on='date').agg({'index' : 'max', 'column1' : 'sum'}))

df

                   index    column1
group   date        
A       2019-01-01  0.0     0.0
        2019-01-03  2.0     2.0
        2019-01-05  4.0     6.0
B       2019-01-02  1.0     1.0
        2019-01-04  3.0     4.0
        2019-01-06  5.0     9.0

Оттуда, если вам нужен формат окончательного вывода без даты, вы можете сделать:

df = df.reset_index().groupby(['group','index'])['column1'].sum()

group  index
A      0.0      0.0
       2.0      2.0
       4.0      6.0
B      1.0      1.0
       3.0      4.0
       5.0      9.0