如何在Polars中过滤后按自定义索引选行（类似Pandas的.loc）

Polars没有Pandas那样内置的持久化行索引，但通过提前将初始行号存储为一列，就能轻松复刻.loc按原索引选行的功能，具体步骤如下：

1. 保留原行索引作为列 #

用.with_row_index()方法为DataFrame添加一列（比如命名为original_idx），用来存储初始的行序号：

import polars as pl

# 创建带原索引列的DataFrame
df = pl.DataFrame({"A": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]}).with_row_index("original_idx")

2. 执行过滤/排序操作 #

后续的过滤、排序等操作都会保留original_idx列，比如模拟你示例中的过滤逻辑：

# 过滤出原索引为0、10、4的行，得到和Pandas中filtered_df等效的结果
filtered_df = df.filter(pl.col("original_idx").is_in([0, 10, 4]))

3. 按原索引检索行 #

通过.filter()结合列匹配，就能像Pandas的.loc一样按原索引选行：

• 检索单行（原索引为4）：

print(filtered_df.filter(pl.col("original_idx") == 4))

输出：

shape: (1, 2)
┌─────────────┬─────┐
│ original_idx ┆ A   │
│ ---         ┆ --- │
│ u32         ┆ i64 │
├─────────────┼─────┤
│ 4           ┆ 4   │
└─────────────┴─────┘

• 检索多行（原索引为0和10）：

print(filtered_df.filter(pl.col("original_idx").is_in([0, 10])))

输出：

shape: (2, 2)
┌─────────────┬─────┐
│ original_idx ┆ A   │
│ ---         ┆ --- │
│ u32         ┆ i64 │
├─────────────┼─────┤
│ 0           ┆ 0   │
│ 10          ┆ 10  │
└─────────────┴─────┘

效率说明 #

这种方案完全基于Polars的向量化操作实现，性能高效，即使处理百万级以上的大型数据集也不会有明显性能损耗。

内容的提问来源于stack exchange，提问作者Míriam