按行索引过滤 LazyFrame

Question

是否有一种惯用的方法从 LazyFrame 获取特定行？我可以想出两种方法。不确定哪个更好，或者我是否应该使用一些不同的方法。

import polars as pl

df = pl.DataFrame({"x": ["a", "b", "c", "d"]}).lazy()

rows = [1, 3]

# method 1
(
  df.with_row_index("row_number")
  .filter(pl.col("row_number").is_in(rows))
  .drop("row_number")
  .collect()
)

# method 2
df.select(pl.all().gather(rows)).collect()

Answer 1

pl.Expr.gather

是按索引获取值的惯用方法。

df.select(pl.all().gather(rows)).collect()

---

为了完整性，可以通过使用索引表达式来改进方法 1。这样就不会创建额外的列。

# method 1.1
df.filter(pl.int_range(pl.len()).is_in(rows)).collect()