Get nested JSON from pandas dataframe grouped by multiple columns

Question

I have a pandas dataframe:

d = {'key': ['foo', 'foo', 'foo', 'foo', 'bar', 'bar', 'bar', 'bar', 'crow', 'crow', 'crow', 'crow'], 
     'date': ['2021-01-01', '2021-01-01', '2021-01-02', '2021-01-02', '2021-01-01', '2021-01-01','2021-01-02', '2021-01-02', '2021-01-01', '2021-01-01', '2021-01-02', '2021-01-02'], 
     'class': [1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2],
     'count': [12, 3, 5, 5, 3, 1, 4, 1, 7, 3, 8, 2],
     'percent': [.8, .2, .5, .5, .75, .25, .8, .2, .7, .3, .8, .2]
}
df = pd.DataFrame(data=d)
df  

     key        date  class  count  percent
0    foo  2021-01-01      1     12     0.80
1    foo  2021-01-01      2      3     0.20
2    foo  2021-01-02      1      5     0.50
3    foo  2021-01-02      2      5     0.50
4    bar  2021-01-01      1      3     0.75
5    bar  2021-01-01      2      1     0.25
6    bar  2021-01-02      1      4     0.80
7    bar  2021-01-02      2      1     0.20
8   crow  2021-01-01      1      7     0.70
9   crow  2021-01-01      2      3     0.30
10  crow  2021-01-02      1      8     0.80
11  crow  2021-01-02      2      2     0.20

I would like to create a nested JSON file that grouped by key and date where count: is a list containing the sums of the counts of key for that day and percent: are lists containing the percentages of the class counts over the total count (there needs to be one list per day containing the percentages of each class).

[
  [
    {
      "key": "foo",
      "count": [
        15,
        10
      ],
      "predictions": [
        [
          .80,
          .20
        ],
        [
          .50,
          .50,
        ]
      ]
    },
    {
      "key": "bar",
      "count": [
        4,
        5
      ],
      "predictions": [
        [
          .75,
          .25
        ],
        [
          .80,
          .20
        ]
      ]
    },
    {
      "key": "crow",
      "count": [
        10,
        10
      ],
      "predictions": [
        [
          .70,
          .30
        ],
        [
          .80,
          .20
        ]
      ]
    }
  ]
]

So far I have:

import json
dfj = dfd.groupby(["key","date"]).apply(lambda x: x.to_dict("r")).to_json(orient="records")
print(json.dumps(json.loads(dfj), indent=2, sort_keys=True))

which returns:

[
  [
    {
      "class": 1,
      "count": 3,
      "date": "2021-01-01",
      "key": "bar",
      "percent": 0.75
    },
    {
      "class": 2,
      "count": 1,
      "date": "2021-01-01",
      "key": "bar",
      "percent": 0.25
    }
  ],
  [
    {
      "class": 1,
      "count": 4,
      "date": "2021-01-02",
      "key": "bar",
      "percent": 0.8
    },
    {
      "class": 2,
      "count": 1,
      "date": "2021-01-02",
      "key": "bar",
      "percent": 0.2
    }
  ],
  [
    {
      "class": 1,
      "count": 7,
      "date": "2021-01-01",
      "key": "crow",
      "percent": 0.7
    },
    {
      "class": 2,
      "count": 3,
      "date": "2021-01-01",
      "key": "crow",
      "percent": 0.3
    }
  ],
  [
    {
      "class": 1,
      "count": 8,
      "date": "2021-01-02",
      "key": "crow",
      "percent": 0.8
    },
    {
      "class": 2,
      "count": 2,
      "date": "2021-01-02",
      "key": "crow",
      "percent": 0.2
    }
  ],
  [
    {
      "class": 1,
      "count": 12,
      "date": "2021-01-01",
      "key": "foo",
      "percent": 0.8
    },
    {
      "class": 2,
      "count": 3,
      "date": "2021-01-01",
      "key": "foo",
      "percent": 0.2
    }
  ],
  [
    {
      "class": 1,
      "count": 5,
      "date": "2021-01-02",
      "key": "foo",
      "percent": 0.5
    },
    {
      "class": 2,
      "count": 5,
      "date": "2021-01-02",
      "key": "foo",
      "percent": 0.5
    }
  ]
]

Any help would be appreciated. Thank you.

Answer 1

You can use:

d   = {'count': ('count', 'sum'), 'predictions': ('percent', list)}
g   = df.groupby(['key', 'date']).agg(**d).groupby(level=0).agg(list)
dct = [{'key': k, **v} for k, v in g.to_dict('i').items()]

Details:

1. groupby the given dataframe on key and date and agg using the dictionary d ,

2. groupby the aggregated frame from step 1 on level=0 and agg using list

3. Finally using to_dict with orient=index to convert the frame from step 2 to dictionary followed by dict comprehension to add the key variable in dictionary.

Result:

[{'key': 'bar', 'count': [4, 5], 'predictions': [[0.75, 0.25], [0.8, 0.2]]},
 {'key': 'crow', 'count': [10, 10], 'predictions': [[0.7, 0.3], [0.8, 0.2]]},
 {'key': 'foo', 'count': [15, 10], 'predictions': [[0.8, 0.2], [0.5, 0.5]]}]