如何使用 Pandas Dataframe 列解析和评估数学表达式？

Question

我想做的是解析这样一个表达式：

result = A + B + sqrt(B + 4)

其中 A 和 B 是数据框的列。所以我必须像这样解析表达式才能得到结果：

new_col = df.B + 4
result = df.A + df.B + new_col.apply(sqrt)

df 是数据框。

我尝试过使用re.sub，但最好只替换列变量（而不是函数），如下所示：

import re

def repl(match):
    inner_word = match.group(1)
    new_var = "df['{}']".format(inner_word)
    return new_var

eq = 'A + 3 / B'
new_eq = re.sub('([a-zA-Z_]+)', repl, eq)
result = eval(new_eq)

所以，我的问题是：

• 是否有 python 库可以做到这一点？如果没有，我怎样才能以简单的方式实现这一目标？
• 创建递归函数可能是解决方案？
• 如果我使用“反向波兰符号”可以简化解析？
• 我必须使用ast 模块吗？

Answer 1

Pandas DataFrames 确实有一个eval 函数。使用您的示例方程式：

import pandas as pd
# create an example DataFrame to work with
df = pd.DataFrame({"A": [1, 2], "B": [3, 4]})
# define equation
eq = 'A + 3 / B'
# actual computation
df.eval(eq)

# more complicated equation
eq = "A + B + sqrt(B + 4)"
df.eval(eq)

警告

请记住，eval 允许运行任意代码，如果您将用户输入传递给此函数，这会使您容易受到代码注入的攻击。

Answer 2

按照@uuazed 提供的示例，更快的方法是使用numexpr

import pandas as pd
import numpy as np
import numexpr as ne

df = pd.DataFrame(np.random.randn(int(1e6), 2), columns=['A', 'B'])
eq = "A + B + sqrt(B + 4)"
timeit df.eval(eq)
# 15.9 ms ± 177 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
timeit A=df.A; B=df.B; ne.evaluate(eq)
# 6.24 ms ± 396 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

numexpr 还可能有更多支持的操作