Move groupby selected columns into a new pandas column as dictionary

user_582175 0
Viewed 340

I have the following pandas DataFrame in python:

df = pd.DataFrame({'id': [1, 1, 2, 2, 3],
                   'field1': [1, 2, 3, 4, 5],
                   'field2': ['a', 'b', 'c', 'd', 'e']})

    id  field1  field2
0   1   1       a
1   1   2       b
2   2   3       c
3   2   4       d
4   3   5       e

I want to group the above table by id then move all selected column values in that group into a new column as a list of python dictionaries.

So from the above I would like to produce this one:

    id  fields
0   1   [{'field1': 1, 'field2': 'a'}, {'field1': 2, 'field2': 'b'}]
2   2   [{'field1': 3, 'field2': 'c'}, {'field1': 4, 'field2': 'd'}]
4   3   [{'field1': 5, 'field2': 'e'}]

I could achieve this with the following python code:

def test(df):
    df['fields'] = [df[['field1', 'field2']].to_dict(orient='records')]*len(df)
    return df

df.groupby('id').apply(test).drop_duplicates('id')[['id', 'fields']]

But I'm sure it can be done better. The question is how? I'm especially dissatisfied with this part:

df['fields'] = [df[['field1', 'field2']].to_dict(orient='records')]*len(df)

where I have to make a list with the length of the groups just to assign the same dictionary value to the rows. Also, this makes it more memory hungry.

pandas dataframe dictionary python
2 Answers

Maybe

df.set_index('id').groupby(level=0).apply(pd.DataFrame.to_dict, orient='r')

id
1    [{'field1': 1, 'field2': 'a'}, {'field1': 2, 'field2': 'b'}]
2    [{'field1': 3, 'field2': 'c'}, {'field1': 4, 'field2': 'd'}]
3    [{'field1': 5, 'field2': 'e'}]
dtype: object

Can always add .to_frame('fields') at the end to get a df back.

user_25356110

Also you can use:

df.groupby('id')['field1','field2'].apply(lambda x: x.to_dict('r')).rename('fields').reset_index()

   id                                             fields
0   1  [{'field1': 1, 'field2': 'a'}, {'field1': 2, '...
1   2  [{'field1': 3, 'field2': 'c'}, {'field1': 4, '...
2   3                     [{'field1': 5, 'field2': 'e'}]
user_118842370
Related