如何按行ID/行号筛选Polars DataFrame？兼谈性能差异

需求场景 #

给定如下Polars DataFrame：

from datetime import datetime

df = pl.DataFrame(
    {
        "integer": [1, 2, 3, 4, 5],
        "date": [
            datetime(2022, 1, 1),
            datetime(2022, 1, 2),
            datetime(2022, 1, 3),
            datetime(2022, 1, 4),
            datetime(2022, 1, 5),
        ],
        "float": [4.0, 5.0, 6.0, 7.0, 8.0],
    }
)

需要根据行号（如[0,4]）提取对应的行子集，得到结果：

shape: (2, 3)
┌─────────┬─────────────────────┬───────┐
│ integer ┆ date                ┆ float │
│ ---     ┆ ---                 ┆ ---   │
│ i64     ┆ datetime[μs]        ┆ f64   │
╞═════════╪═════════════════════╪═══════╡
│ 1       ┆ 2022-01-01 00:00:00 ┆ 4.0   │
│ 5       ┆ 2022-01-05 00:00:00 ┆ 8.0   │
└─────────┴─────────────────────┴───────┘

可行实现方式 #

1. 方括号索引（推荐） #

直接使用方括号传入行号列表即可实现类似take的功能：

take_indices = [0, 4]
result = df[[take_indices]]
print(result)

该方法语法简洁，性能优于列级take操作。

2. 列级take（不推荐） #

你尝试过的df.select(pl.all().take([take_indices]))也能实现需求，但性能较差，不适用于超大规模数据集（如1亿行以上）。

性能疑问解答：为什么filter比take/方括号索引更快？ #

在5000万行数据集的测试中，filter耗时0.18秒，方括号索引0.32秒，select.gather0.5秒，核心原因在于两者底层实现逻辑的差异：

• filter基于布尔位图优化：当使用pl.int_range(0, pl.count()).is_in(take_indices)作为筛选条件时，Polars会将is_in操作转换为高效的位图运算。位图可以批量标记需要保留的行，CPU缓存局部性更好，能快速跳过不需要的行，尤其当目标行占比极低时，这种优势更明显。
• take/方括号索引基于随机访问：这类方法需要根据给定行号逐个定位数据，当行号离散分布时，内存访问模式是随机的，无法充分利用CPU缓存的连续读取优化，导致整体效率低于位图筛选。

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