如何用Ibis筛选出每个国家中人口最多的城市？

原始数据 #

┏━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ 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 实现方式 #

方法一：窗口函数法 #

通过窗口函数给每个国家内的城市按人口降序排名，筛选排名第一的行即可：

import ibis

# 准备数据
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]
}

# 构建Ibis内存表
con = ibis.memtable(data)
table = con.table

# 定义窗口：按country分组，population降序排列
rank_window = ibis.window(by='country', order_by=ibis.desc('population'))

# 添加组内排名，筛选排名第一的行
result = table.mutate(rank=ibis.row_number().over(rank_window)).filter(ibis._.rank == 1).drop('rank')

# 执行查询并输出结果
print(result.execute())

方法二：分组聚合关联法 #

先计算每个国家的最大人口，再关联原表筛选匹配行：

import ibis

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]
}

con = ibis.memtable(data)
table = con.table

# 分组计算每个国家的最大人口
max_pop_table = table.group_by('country').agg(max_pop=table.population.max())

# 关联原表，筛选人口等于对应国家最大人口的行
result = table.join(max_pop_table, on='country').filter(table.population == max_pop_table.max_pop).drop('max_pop')

print(result.execute())

内容的提问来源于stack exchange，提问作者ianmcook