使用列值对迭代器从 Polars DataFrame 检索行的简洁方法

Question

我经常需要从给定列值集合的 Polars DataFrame 中检索行，就像我可能在数据库中使用复合键一样。这在 Polars 中可以使用

DataFrame.row

实现，但生成的表达式非常冗长：

row_index = {'treatment': 'red', 'batch': 'C', 'unit': 76}

row = df.row(by_predicate=(
    (pl.col('treatment') == row_index['treatment'])
    & (pl.col('batch') == row_index['batch'])
    & (pl.col('unit') == row_index['unit'])
))

我发现的最简洁的方法是

from functools import reduce
from operator import and_

expr = reduce(and_, (pl.col(k) == v for k, v in row_index.items()))

row = df.row(by_predicate=expr)

但这仍然冗长且难以阅读。有更容易的方法吗？可能是我缺少内置的 Polars 功能？

Answer 1

如果所有条件都为 true，

(a == b) & (c == d)

将返回 true。

另一种表达方式是用

pl.all_horizontal()

pl.all_horizontal(a == b, c == d)

• pl.any_horizontal()

  可用于“逻辑或”

传递多个条件的一种方法是使用理解：

predicate = pl.all_horizontal(
   pl.col(name) == value 
   for name, value in row_index.items()
)

df.row(by_predicate=predicate)