Group rows by certain timeperiod dending on other factors

Question

What I start with is a large dataframe (more than a million entires) of this structure:

id  datetime             indicator  other_values ...
1   2020-01-14 00:12:00  0          ...
1   2020-01-17 00:23:00  1          ...
  ... 
1   2021-02-01 00:00:00  0          ...
2   2020-01-15 00:05:00  0          ...
2   2020-03-10 00:07:00  0          ...
  ... 
2   2021-05-22 00:00:00  1          ...
  ... 

There is no specific order other than a sort by id and then datetime. The dataset is not complete (there is not data for every day, but there can be multiple entires of the same day).

Now for each time where indicator==1 I want to collect every row with the same id and a datetime that is at most 10 days before. All other rows which are not in range of the indicator can be dropped. In the best case I want it to be saved as a dataset of time series which each will be later used in a Neural network. (There can be more than one indicator==1 case per id, other values should be saved).

An example for one id: I want to convert this

id  datetime             indicator  other_values ...
1   2020-01-14 00:12:00  0          ...
1   2020-01-17 00:23:00  1          ...
1   2020-01-17 00:13:00  0          ...
1   2020-01-20 00:05:00  0          ...
1   2020-03-10 00:07:00  0          ...
1   2020-05-19 00:00:00  0          ...
1   2020-05-20 00:00:00  1          ...

into this

id  datetime             group      other_values ...
1   2020-01-14 00:12:00  A          ...
1   2020-01-17 00:23:00  A          ...
1   2020-01-17 00:13:00  A          ...
1   2020-05-19 00:00:00  B          ...
1   2020-05-20 00:00:00  B          ...

or a similar way to group into group A, B, ... .

A naive python for-loop is not possible due to taking ages for a dataset like this. There is propably a clever way to use df.groupby('id'), df.groupby('id').agg(...), df.sort_values(...) or df.apply(), but I just do not see it.

Answer 1

I'm not aware of a way to do this with df.agg, but you can put your for loop inside the groupby using .apply(). That way, your comparisons/lookups can be done on smaller tables, then groupby will handle the re-concatenation:

import pandas as pd
import datetime
import uuid

df = pd.DataFrame({
    "id": [1, 1, 1, 2, 2, 2],
    "datetime": [
        '2020-01-14 00:12:00',
        '2020-01-17 00:23:00',
        '2021-02-01 00:00:00',
        '2020-01-15 00:05:00',
        '2020-03-10 00:07:00',
        '2021-05-22 00:00:00',
    ],
    "indicator": [0, 1, 0, 0, 0, 1]
})
df.datetime = pd.to_datetime(df.datetime)

timedelta = datetime.timedelta(days=10)
def consolidate(grp):
    grp['Group'] = None
    for time in grp[grp.indicator == 1]['datetime']:
        grp['Group'][grp['datetime'].between(time - timedelta, time)] = uuid.uuid4()
    return grp.dropna(subset=['Group'])

df.groupby('id').apply(consolidate)

If there are multiple rows with indicator == 1 in each id grouping, then the for loop will apply in index order (so a later group might overwrite an earlier group). If you can be certain that there is only one indicator == 1 in each grouping, we can simplify the consolidate function:

def consolidate(grp):
    time = grp[grp.indicator == 1]['datetime'].iloc[0]
    grp = grp[grp['datetime'].between(time - timedelta, time)]
    grp['Group'] = uuid.uuid4()
    return grp