我可以使用类似 pandas 的字符串表达式来过滤 DataFrame 吗？

Question

我正在考虑在一个允许用户输入谓词表达式以过滤/子集数据行的工具中用

pandas

替换我对

polars

的使用。这允许用户使用

pandas.DataFrame.query

方法可以解析的表达式，例如

"x > 1"

，作为一个非常简单的示例。

但是，我似乎找不到一种方法来使用与

polars.DataFrame.filter

相同类型的字符串表达式，以便我可以将

pandas

替换为

polars

，而不需要用户更改其谓词表达式。

我发现的唯一接近我的问题的是这个帖子：String as a condition in a filter

不幸的是，这并不是我所需要的，因为它仍然需要像

"pl.col('x') > 1"

这样的字符串表达式，而不是简单地

"x > 1"

。

有没有办法在

polars

中使用更简单（“不可知”）的语法？

使用

polars.DataFrame.filter

文档中的示例：

>>> df = pl.DataFrame(
...     {
...         "foo": [1, 2, 3],
...         "bar": [6, 7, 8],
...         "ham": ["a", "b", "c"],
...     }
... )

打电话给

df.filter

时，我被迫使用如下的表达方式：

pl.col("foo") < 3
(pl.col("foo") < 3) & (pl.col("ham") == "a")

但是，我希望能够分别使用以下字符串表达式，以便该工具的用户（当前使用

pandas

）不必了解

polars

特定的语法（因此允许我在不影响用户的情况下交换库）：

"foo < 3"
"foo < 3 & ham == 'a'"

当我尝试这样做时，会发生以下情况，这令人费解，因为

str

是谓词参数支持的类型之一，因此不清楚

str

谓词支持的语法，因为文档没有显示任何这样的例子：

>>> df.filter("foo < 3")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/Caskroom/miniconda/base/envs/gedi_subset/lib/python3.10/site-packages/polars/internals/dataframe/frame.py", line 2565, in filter
    self.lazy()
  File "/usr/local/Caskroom/miniconda/base/envs/gedi_subset/lib/python3.10/site-packages/polars/utils.py", line 391, in wrapper
    return fn(*args, **kwargs)
  File "/usr/local/Caskroom/miniconda/base/envs/gedi_subset/lib/python3.10/site-packages/polars/internals/lazyframe/frame.py", line 1165, in collect
    return pli.wrap_df(ldf.collect())
exceptions.NotFoundError: foo < 3

我期待的是与

df.filter(pl.col("foo") < 3)

返回相同的返回值。

Answer 1

您可以尝试使用

pl.sql

来实现这一点。

import polars as pl

df = pl.DataFrame(
    {
        "foo": [1, 2, 3],
        "bar": [6, 7, 8],
        "ham": ["a", "b", "c"],
    }
)

string_expr = "foo < 3 and ham = 'a'"

pl.sql(f"""
SELECT * FROM df
WHERE {string_expr}
""").collect()

shape: (1, 3)
┌─────┬─────┬─────┐
│ foo ┆ bar ┆ ham │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ str │
╞═════╪═════╪═════╡
│ 1   ┆ 6   ┆ a   │
└─────┴─────┴─────┘

请注意，

SQL

语言不像 pandas 那样使用按位

&

或相等

==

，因此您可能需要将

&

替换为

and

，将

==

替换为

=

。