[Python] What's the best way to group by more than one column?

Essentially, I want to group all records that match either key

Input:

|id|key1|key2|
|-|-|-|
|1|a|x|
|2|a|y|
|3|b|y|
|4|c|z|


Desired output: (The first 3 rows are grouped together)

|key1 (or any identifier, I only care about the final counts)|len|
|-|-|
|a|3|
|c|1|


My attempted (incomplete) solution:

import polars as pl

df = pl.DataFrame({
'id': [1, 2, 3, 4],
'key1': ['a', 'a', 'b', 'c'],
'key2': ['x', 'y', 'y', 'z'],
})
df = (
df.group_by(
'key1',
'key2',
)
.len()
.group_by('key1')
.agg(
pl.col('key2'),
pl.col('len').sum()
)
)
print(df)


Which gives:

|key1|key2|len|
|-|-|-|
|a|["y", "x"]|2|
|b|["y"]|1|
|c|["z"]|1|


However, I'm not sure how I would further group this by key2 (merging the a and b rows, since they have a common value of y) while preserving the sum of len

Continue reading...