如何按row_id/row_number过滤数据帧

我希望基于类似于的数据帧的row_id/row_number获得行的子集 pyarrow.Table.take .例如:给定以下数据帧

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

print(df)

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

我正在寻找类似take的函数df.take([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   │
└─────────┴─────────────────────┴───────┘

行号是其他过程的结果,并已移交。尝试使用 df.select(pl.all().take([take_indices]) 并注意到它比实际直接运行过滤器慢。即 df.filter(filter_expr) 。请注意,我是在超大数据集(>100m行)上进行此操作的。

编辑:谢谢你的回答。使用df[[take_dices]]有效。然而,仍然很奇怪的是,为什么过滤器仍然优于select.gather和方括号方法。我的数据集上有50米行的计时:

select.gather:.5s 方括号:.32s[与mozway的计时一致] 过滤器:.18s