使用 Ibis 过滤表格以筛选出每组中具有最大值的行

Question

我有一张这样的桌子：

┏━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ 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 │
└───────────────┴─────────────┴────────────┘

我想过滤这个表，只返回每个国家人口最多的城市。所以结果应该是这样的（行的顺序并不重要）：

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

有了pandas，我可以这样做：

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]

如何使用 Ibis 做到这一点？

Answer 1

有了宜必思，你可以这样做：

import ibis
from ibis import _
import ibis.selectors as s

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)
     .select(~s.c('row_num'))
)

它的作用：

1. 添加一列

   row_num

   ，按人口（从最大到最小）对每个国家/地区的城市进行排名。
2. 筛选至每个国家/地区排名 0（最大）的城市。
3. 删除

   row_num

   列（使用 Ibis 选择器）。

这使用 Ibis 的 underscore API 来简化链接。