如何检查用户输入是否存在于DataFrame列值中并返回匹配行数据

场景与基础信息 #

• 需求：校验命令行传入的用户输入值是否存在于数据库查询生成的DataFrame指定列中，匹配成功返回对应行的关联字段值
• 示例SQL：

my_query = "SELECT a.fruits, a.colors FROM table1"

• 原始数据库表table1数据：

• 查询生成的DataFrame结构：

• 已完成的基础代码：

cursor = conn.cursor()
cursor.execute(my_query)
df = pd.DataFrame(cursor.fetchall(), columns=['fruits','colors'])

parser = argparse.ArgumentParser()
group = parser.add_mutually_exclusive_group(required=True)
group.add_argument('-a','--add')
args = parser.parse_args()
user_input = args.add.lower()
conn.close()

• 测试执行命令：

python3 your_script.py -a grapes

• 匹配成功预期输出：

grapes, purple

实现代码 #

直接在现有代码后追加如下逻辑即可：

# 忽略大小写匹配fruits列的对应行，避免大小写差异导致匹配失败
match_result = df[df['fruits'].str.lower() == user_input]

if not match_result.empty:
    # 取第一条匹配数据，若允许多个重复值可遍历match_result输出所有结果
    target_fruit = match_result.iloc[0]['fruits']
    target_color = match_result.iloc[0]['colors']
    print(f"{target_fruit}, {target_color}")
else:
    # 匹配失败的提示逻辑可按需调整
    print(f"未找到与输入「{user_input}」匹配的记录")

说明 #

• 对DataFrame的fruits列统一转小写后再和用户输入比对，可以兼容用户输入大小写不规范的场景，比如输入Grapes/GRAPES都能正常匹配
• 若业务上fruits列存在重复值，可循环遍历match_result逐行输出，不需要只取iloc[0]

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