Group by and aggregate rows into list of series or dicts in Pandas

Viewed 181

I have a dataframe with 7 million rows which looks like this

|    | ID     | VAL1 | VAL2
|---:|:-------|:-----|:----
|  0 | QWERTY | 1    | ABC
|  1 | 123456 | 2    | ABC
|  2 | QWERTY | 3    | ABC
|  3 | QWERTY | 4    | ABC
|  4 | 123456 | 5    | ABC
df = pd.DataFrame(
    columns=("ID", "VAL1", "VAl2"),
    data=[("QWERTY", 1, "ABC"), ("123456", 2, "ABC"), ("QWERTY", 3, "ABC"), ("QWERTY", 4, "ABC"), ("123456", 5, "ABC")]
)

And I want to groupby it by ID or transform it in a shape like this

|    | ID     | GROUPED
|---:|:-------|:------------------------------------------------
|  0 | QWERTY | [{'ID': 'QWERTY', 'VAL1': 1, 'VAl2': 'ABC'}, {'ID': 'QWERTY', 'VAL1': 3, 'VAl2': 'ABC'}, {'ID': 'QWERTY', 'VAL1': 4, 'VAl2': 'ABC'}, ]
|  1 | 123456 | [{'ID': '123456', 'VAL1': 2, 'VAl2': 'ABC'}, {'ID': '123456', 'VAL1': 5, 'VAl2': 'ABC'}]

It should be grouped by ID and have a list with whole rows which corresponds to that id. Rows in the list can be either Series or Dict.


I've tried to do that in this way

test1 = df.groupby("ID").apply(lambda x: df.iloc[list(x.to_dict()["VAL1"].keys())])

But it expands the rows and creates multi-index, instead of a list or something

              ID  VAL1 VAl2
ID
123456 1  123456     2  ABC
       4  123456     5  ABC
QWERTY 0  QWERTY     1  ABC
       2  QWERTY     3  ABC
       3  QWERTY     4  ABC

Is there any way to do it with pandas?
Unfortunately, plain python loops/maps are quite slow when operating with 5 million rows.

2 Answers

Try this

compact_df =  df.groupby('ID').apply(lambda group: group.to_dict(orient='records'))

Try this:

test1 = df.groupby("ID").apply(lambda x: x.to_dict(orient='records'))

res=pd.DataFrame(test1, index=test1.index, columns=['GROUPED'])

>>> print(res)
                                                  GROUPED
ID
123456  [{'ID': '123456', 'VAL1': 2, 'VAl2': 'ABC'}, {...
QWERTY  [{'ID': 'QWERTY', 'VAL1': 1, 'VAl2': 'ABC'}, {...
Related