Streamlit中Pandas merge函数未按预期工作问题求助

问题概述 #

在Streamlit的回调函数中使用Pandas的merge方法时，DataFrame无法正确更新，但drop方法可正常工作。无论是否启用st.cache装饰器，该问题均会出现。

复现步骤 #

问题代码片段 #

import streamlit as st
import pandas as pd


@st.cache(allow_output_mutation=True)
def read_df():
    df = pd.DataFrame({
        'col1':[1,2],
        'col2':['A','B']
    })
    return df

df = read_df()

def do_something():
    global df
    df_new = pd.DataFrame({
        'col1':[1,2],
        'col3':["X","Y"]
    })
    df.drop(['col2'], axis = 1, inplace = True)
    df = df.merge(df_new, on="col1")

st.button("Do Something", on_click=do_something, args =())

download_csv = df.to_csv().encode('utf-8')
st.download_button('Download', data = download_csv, file_name = 'download_csv.csv', mime='text/csv')

操作步骤 #

1. 点击页面上的「Do Something」按钮
2. 点击「Download」按钮下载CSV文件

预期结果 #

下载的CSV内容应为：

,col1,col3
0,1,X
1,2,Y

实际结果 #

实际下载的CSV内容缺失col3列：

,col1
0,1
1,2

调试环境信息 #

• Streamlit版本：1.16.0
• Python版本：3.8.15
• 环境：Conda
• 操作系统：Windows 11
• 浏览器：Edge v108.0.1462.54

解决方案 #

问题根源在于使用全局变量存储DataFrame，而Streamlit每次交互都会重新运行整个脚本，全局变量会被重置为初始值。改用st.session_state来持久化DataFrame状态即可解决：

修改后的代码 #

import streamlit as st
import pandas as pd


@st.cache(allow_output_mutation=True)
def read_df():
    df = pd.DataFrame({
        'col1':[1,2],
        'col2':['A','B']
    })
    return df

# 初始化session_state中的DataFrame
if 'df' not in st.session_state:
    st.session_state.df = read_df()

def do_something():
    df_new = pd.DataFrame({
        'col1':[1,2],
        'col3':["X","Y"]
    })
    # 操作session_state中的DataFrame
    st.session_state.df.drop(['col2'], axis=1, inplace=True)
    st.session_state.df = st.session_state.df.merge(df_new, on="col1")

st.button("Do Something", on_click=do_something)

# 使用session_state中的DataFrame生成下载内容
download_csv = st.session_state.df.to_csv().encode('utf-8')
st.download_button('Download', data=download_csv, file_name='download_csv.csv', mime='text/csv')

原理说明 #

st.session_state是Streamlit提供的状态管理工具，能在用户交互之间保留变量值，避免每次脚本重运行时重置数据。通过将DataFrame存储在session_state中，回调函数的修改会被持久化，下载操作也能获取到更新后的DataFrame。

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