Pandas query后用assign结合字典推导式计算异常的解决问询

场景与问题 #

想要通过Pandas链式调用实现：先使用query筛选数据，再通过assign结合字典推导式批量生成加权列，但运行代码后结果不符合预期。

示例代码：

import pandas as pd

df = pd.DataFrame({
    "a": [1, 2, 3],
    "b": [4, 5, 6],
    "weight": [0.1, 0.2, 0.3]
})

metrics = ["a", "b"]

df = df.query("b > 4").assign(
    **{
        f"weighted_{metric}": lambda df: df[metric] * df["weight"]
        for metric in metrics
    }
)
print(df)

错误输出：

a  b  weight  weighted_a  weighted_b
1  2  5     0.2         1.0         1.0
2  3  6     0.3         1.8         1.8

预期weighted_a应为0.4、0.9，实际结果与weighted_b完全一致，显然错误。

原因分析 #

字典推导式中的lambda表达式存在变量捕获延迟问题：循环过程中，lambda并不会立即绑定当前的metric值，而是在实际执行时才去读取metric的当前值。当循环结束后，所有lambda里的metric都指向循环的最后一个元素（即"b"），因此所有加权列都在计算b * weight，导致结果错误。

解决方法（适配链式调用+支持复杂计算） #

方法1：用闭包固定变量值 #

通过一个外层函数将每个循环的metric值固定，返回绑定了该值的lambda，避免延迟捕获问题。这种方式可以轻松扩展到复杂计算逻辑。

修正后的代码：

import pandas as pd

df = pd.DataFrame({
    "a": [1, 2, 3],
    "b": [4, 5, 6],
    "weight": [0.1, 0.2, 0.3]
})

metrics = ["a", "b"]

# 闭包函数：固定当前循环的metric值，返回对应计算逻辑
def build_weighted_calculator(metric):
    return lambda df: df[metric] * df["weight"]

df = df.query("b > 4").assign(
    **{
        f"weighted_{metric}": build_weighted_calculator(metric)
        for metric in metrics
    }
)
print(df)

正确输出：

a  b  weight  weighted_a  weighted_b
1  2  5     0.2         0.4         1.0
2  3  6     0.3         0.9         1.8

扩展：适配复杂计算函数 #

如果实际场景需要更复杂的计算逻辑，只需修改闭包内的计算代码即可：

def build_weighted_calculator(metric):
    def complex_calculation(df):
        # 示例复杂逻辑：加权值 + 该列的均值
        weighted_val = df[metric] * df["weight"]
        return weighted_val + df[metric].mean()
    return complex_calculation

此方法完全兼容链式调用结构，同时支持任意复杂的自定义计算。

内容的提问来源于stack exchange，提问作者David Masip