Use Ibis to filter table to row with largest value in each group

Question

I have a table like this:

┏━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ country       ┃ city        ┃ population ┃
┡━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━━┩
│ string        │ string      │ int64      │
├───────────────┼─────────────┼────────────┤
│ India         │ Bangalore   │    8443675 │
│ India         │ Delhi       │   11034555 │
│ India         │ Mumbai      │   12442373 │
│ United States │ Los Angeles │    3820914 │
│ United States │ New York    │    8258035 │
│ United States │ Chicago     │    2664452 │
│ China         │ Shanghai    │   24281400 │
│ China         │ Guangzhou   │   13858700 │
│ China         │ Beijing     │   19164000 │
└───────────────┴─────────────┴────────────┘

I want to filter this table, returning only the most populous city in each country. So the result should look like this (order of rows does not matter):

┏━━━━━━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ country       ┃ city     ┃ population ┃
┡━━━━━━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━━━┩
│ string        │ string   │ int64      │
├───────────────┼──────────┼────────────┤
│ India         │ Mumbai   │   12442373 │
│ United States │ New York │    8258035 │
│ China         │ Shanghai │   24281400 │
└───────────────┴──────────┴────────────┘

With pandas, I can do it like this:

import pandas as pd

df = pd.DataFrame(data={'country': ['India', 'India', 'India', 'United States', 'United States', 'United States', 'China', 'China', 'China'],
                        'city': ['Bangalore', 'Delhi', 'Mumbai', 'Los Angeles', 'New York', 'Chicago', 'Shanghai', 'Guangzhou', 'Beijing'],
                        'population': [8443675, 11034555, 12442373, 3820914, 8258035, 2664452, 24281400, 13858700, 19164000]})

idx = df.groupby('country').population.idxmax()
df.loc[idx]

How do I do this with Ibis?

Answer 1

With Ibis, you can do it like this (run this code after the code in the question):

import ibis
from ibis import _

ibis.options.interactive = True

t = ibis.memtable(df)

(
    t.mutate(row_num=ibis.row_number().over(group_by=_.country, order_by=_.population.desc()))
     .filter(_.row_num==0)
     .drop('row_num')
)

What this does:

1. Adds a column row_num that ranks the cities in each country by population (from largest to smallest).
2. Filters down to the rank 0 (largest) city in each country.
3. Removes the row_num column.

This uses Ibis's underscore API to simplify chaining.