Python如何创建函数，根据指定列的值返回特定列数据？

问题场景 #

给定如下示例DataFrame：

one  two  three  year
0   12   34    561  2001
1    2   67    781  1999
2   34   12     90  2001
3    5   67     89  2011

想要实现一个函数，传入年份后返回one和two两列的对应数据，期望输出（以2001年为例）：

one  two
0   12   34
2   34   12

当前编写的代码调用后仅返回列名，无数据：

def one_two(year):
 return df.loc[df['year'] == year,['one','two']].sort_values(by='one')

# 调用
one_two(2001)
# 输出仅列名：
# one  two

问题原因及解决方案 #

核心原因 #

1. 数据类型不匹配：year列的数据类型（如字符串）与传入的参数类型（如整数）不一致，导致匹配条件df['year'] == year返回全False，最终得到空DataFrame。
2. 变量作用域问题：函数未正确访问到目标DataFramedf（比如df定义在其他局部作用域中，函数无法读取）。

修正方案 #

方案1：将DataFrame作为参数传入（推荐，代码更健壮）

把df作为函数参数传入，同时添加类型匹配检查：

import pandas as pd

# 先构造你的DataFrame
data = {
    'one': [12, 2, 34, 5],
    'two': [34, 67, 12, 67],
    'three': [561, 781, 90, 89],
    'year': [2001, 1999, 2001, 2011]
}
df = pd.DataFrame(data)

def one_two(df, year):
    # 检查year列类型，若为字符串则转换传入参数类型
    if df['year'].dtype == object:
        year = str(year)
    # 筛选并排序
    return df.loc[df['year'] == year, ['one', 'two']].sort_values(by='one')

# 调用示例
print(one_two(df, 2001))

方案2：确保全局变量可访问（适用于简单场景）

如果df是全局变量，确保函数能读取到，同时校验类型：

import pandas as pd

# 全局定义DataFrame
data = {
    'one': [12, 2, 34, 5],
    'two': [34, 67, 12, 67],
    'three': [561, 781, 90, 89],
    'year': [2001, 1999, 2001, 2011]
}
df = pd.DataFrame(data)

def one_two(year):
    # 类型匹配处理
    if isinstance(df['year'].iloc[0], str):
        year = str(year)
    return df.loc[df['year'] == year, ['one', 'two']].sort_values(by='one')

# 调用
print(one_two(2001))

验证结果 #

执行修正后的代码，调用one_two(2001)会返回期望的结果：

one  two
0   12   34
2   34   12

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