groupby to return nth group - NOT row

Viewed 861

I'm attempting to group by two factors in a long (>2M) rows.

Background to the data

The second factor is effectively a test date - for a given sample (the first group) a sample can be retested. However the test can change the sample, so it useful to be able to selectively pull out the batch of test data from the second test (or third, fourth, etc.)

Earlier question and current code

I asked a question sometime ago, and got somewhere with a function, thanks very much to Praveen: Sorting and ranking by dates, on a group in a pandas df but this is incredibly slow: running on a 1.2M row dataset, it has so far taken over 20 mins on a JupyterHub server with 96gb ram, and is still going.

df1 = df.sort_values(['id', 'date'], ascending=[True, False])
df1['date_rank'] = df1.groupby(['id'])['date'].transform(lambda x: list(map(lambda y: dict(map(reversed, dict(enumerate(x.unique())).items()))[y]+1,x)) )

Current idea

Instead of labelling the test dates with a number, as was my first idea, it feels like this is close to achievable with cumcount, groupby.nth or getgroup - but I can't get there. Anybody have any ideas? The dataset below is still relevant:

df = pd.DataFrame({
'id':[1, 1, 1, 2, 2, 3, 3, 3, 4, 4, 5, 5, 6,6,6,7,7], 
'value':[.01, .4, .5, .7, .77, .1,.2, 0.3, .11, .21, .4, .01, 3, .5, .8, .9, .1],
'date':['10/01/2017 15:45:00','10/01/2017 15:45:00','05/01/2017 15:56:00',
    '11/01/2017 15:22:00','11/01/2017 15:22:00','06/01/2017 11:02:00','05/01/2017 09:37:00','05/01/2017 09:37:00','05/01/2017 09:55:00',
    '05/01/2017 09:55:00','05/01/2017 10:08:00','05/01/2017 10:09:00','03/02/2017 08:55:00',
    '03/02/2017 09:15:00','03/02/2017 09:31:00','09/01/2017 15:42:00',
    '19/01/2017 16:34:00']})

And i'm trying to generate data like that in date_rank column; although i'm happy to not generate a column, i just want to be able to select rows containing the second test date data.

enter image description here

For example:

df.groupby(['id','date'])['value'].unique()

Returns:

id  date               
1   05/01/2017 15:56:00           [0.5]
    10/01/2017 15:45:00     [0.01, 0.4]
2   11/01/2017 15:22:00     [0.7, 0.77]
3   05/01/2017 09:37:00      [0.2, 0.3]
    06/01/2017 11:02:00           [0.1]
4   05/01/2017 09:55:00    [0.11, 0.21]
5   05/01/2017 10:08:00           [0.4]
    05/01/2017 10:09:00          [0.01]
6   03/02/2017 08:55:00           [3.0]
    03/02/2017 09:15:00           [0.5]
    03/02/2017 09:31:00           [0.8]
7   09/01/2017 15:42:00           [0.9]
    19/01/2017 16:34:00           [0.1]

I would want to be able to select (for example) the second group of a group where available e.g.:

1 - 10/01/2017 15:45:00 [0.01, 0.04]
3 - 06/01/2017 11:02:00 [0.1]
5 - 05/01/2017 10:09:00 [0.01]
6 - 03/02/2017 09:15:00 [0.5]

Any help or thoughts greatly appreciated.

1 Answers

You can reset_index, then use GroupBy + nth:

res = df.reset_index().groupby('id').nth(1)

print(res)

                   date        value
id                                  
1   10/01/2017 15:45:00  [0.01, 0.4]
3   06/01/2017 11:02:00        [0.1]
5   05/01/2017 10:09:00       [0.01]
6   03/02/2017 09:15:00        [0.5]
7   19/01/2017 16:34:00        [0.1]
Related