如何在函数中使用df.query方法实现DataFrame多条件过滤？

我需要创建一个函数，通过df.query方法根据多个条件过滤DataFrame，最终返回包含符合条件的Value列数据的Series。

示例DataFrame #

我已经能手动写出过滤代码：

test = df.query('Week == 1 & Treatment.str.startswith("A").values')['Value']

但不想每次换Week或Treatment条件都重复复制这段代码，所以想把它封装成函数，将目标Week和Treatment首字母作为参数传入，返回对应的数据。

我写了初始函数，但参数传不进去：

def sorter(df,week):
    return df.query('Week == week & Treatment.str.startswith("A").values')['Value']

运行时会报错week未定义，不知道怎么把函数的week参数传入df.query的字符串条件里。

解决方法 #

方法1：使用@符号引用外部变量（推荐）

pandas的query方法支持用@符号直接引用函数中的外部参数，这是最安全且简洁的方式：

def sorter(df, week):
    return df.query('Week == @week & Treatment.str.startswith("A").values')['Value']

方法2：用f-string格式化条件字符串

通过字符串格式化把参数值直接插入到查询字符串中：

def sorter(df, week):
    return df.query(f'Week == {week} & Treatment.str.startswith("A").values')['Value']

扩展：让Treatment前缀也支持参数

如果想让Treatment的过滤条件也动态可调，可以把前缀也做成参数：

def sorter(df, week, treatment_prefix):
    return df.query('Week == @week & Treatment.str.startswith(@treatment_prefix).values')['Value']

调用示例：

# 获取Week=1、Treatment以"A"开头的Value列数据
result = sorter(df, 1, "A")

注意：你的代码中之前写的是['value']，但DataFrame的列名是Value（大写V），需要保持一致，否则会报错找不到列。

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