Streamlit中groupby.transform填充Test/Val数据无效问题排查

我用以下代码通过groupby.transform按分组均值填充train、test、val数据集的缺失值，Jupyter Notebook里运行正常，所有数据集缺失值都能正确填充，但在Streamlit里只有train数据集的填充生效，test和val的缺失值没被填充，求原因？

相关代码 #

示例测试代码 #

import pandas as pd
import numpy as np

df1 = pd.DataFrame({'value': [1, np.nan, np.nan, 2, 3, 1, 3, np.nan, 3], 'name': ['A','A', 'B','B','B','B', 'C','C','C']})
df2 = pd.DataFrame({'value': [1, np.nan, np.nan, 2, 3, 1, 3, np.nan, 3], 'name': ['A','A', 'B','B','B','B', 'C','C','C']})

# df1输出
#   name  value
# 0    A      1
# 1    A    NaN
# 2    B    NaN
# 3    B      2
# 4    B      3
# 5    B      1
# 6    C      3
# 7    C    NaN
# 8    C      3

# df2输出
#   name  value
# 0    A      1
# 1    A    NaN
# 2    B    NaN
# 3    B      2
# 4    B      3
# 5    B      1
# 6    C      3
# 7    C    NaN

numeric_only_columns = df1.select_dtypes(exclude = ['object', 'datetime']).columns.to_list()

# 假设groupby_columns为['name']
groupby_columns = ['name']
for i in numeric_only_columns:
    df1[i] = df1[i].fillna(df1.groupby(groupby_columns)[i].transform('mean'))
    df2[i] = df2[i].fillna(df1.groupby(groupby_columns)[i].transform('mean'))

Streamlit代码 #

import streamlit as st
import pandas as pd

# 假设train、test、val已提前加载
fill_na_columns = st.selectbox("select method?",('mean','min','max','median'))
groupby_columns = st.multiselect("Groupby columns select", train.columns.to_list(), default=train.columns.to_list()[0])

if fill_na_columns == 'mean':
   st.dataframe(train.groupby(groupby_columns).mean(numeric_only=True))
   numeric_only_columns = train.select_dtypes(exclude = ['object', 'datetime']).columns.to_list()

   for i in numeric_only_columns:
       train[i] = train[i].fillna(train.groupby(groupby_columns)[i].transform('mean'))
       test[i] = test[i].fillna(train.groupby(groupby_columns)[i].transform('mean'))
       val[i] = val[i].fillna(train.groupby(groupby_columns)[i].transform('mean'))

原因分析 #

1. Streamlit重运行机制导致数据重置：Streamlit每次用户交互（选择填充方法、分组列）都会重新执行整个脚本。如果train、test、val是在脚本开头直接加载的原始数据，每次重运行都会把它们重置回初始状态——看起来test/val没被填充，其实是填充后又被原始数据覆盖了。
2. transform的索引对齐问题：train.groupby(...)['mean']返回的是和train行数匹配的序列，直接赋值给test/val时，pandas会按索引对齐。如果test/val的索引和train不匹配，只有索引重叠的行能被填充，其余行仍保留NaN。
3. 分组键不匹配：如果test/val中存在train里没有的分组，transform无法找到对应均值，对应位置会保留NaN，看起来像是没填充。

解决办法 #

1. 用Session State保存数据集状态 #

把train、test、val存在Streamlit的st.session_state中，避免每次重运行重置：

import streamlit as st
import pandas as pd

# 初始化Session State，仅第一次运行时加载原始数据
if "train" not in st.session_state:
    st.session_state.train = pd.read_csv("train.csv")  # 替换为你的数据加载逻辑
if "test" not in st.session_state:
    st.session_state.test = pd.read_csv("test.csv")
if "val" not in st.session_state:
    st.session_state.val = pd.read_csv("val.csv")

# 后续操作基于session_state中的数据集
fill_na_columns = st.selectbox("选择填充方法?", ('mean','min','max','median'))
groupby_columns = st.multiselect("选择分组列", 
                                 st.session_state.train.columns.to_list(), 
                                 default=st.session_state.train.columns.to_list()[0])

if fill_na_columns == 'mean':
    st.dataframe(st.session_state.train.groupby(groupby_columns).mean(numeric_only=True))
    numeric_only_columns = st.session_state.train.select_dtypes(exclude=['object', 'datetime']).columns.to_list()

    for i in numeric_only_columns:
        # 填充train
        st.session_state.train[i] = st.session_state.train[i].fillna(
            st.session_state.train.groupby(groupby_columns)[i].transform('mean')
        )
        # 计算train的分组均值，用merge映射到test/val，避免索引对齐问题
        group_means = st.session_state.train.groupby(groupby_columns)[i].mean().reset_index()
        # 填充test
        merged_test = st.session_state.test.merge(group_means, on=groupby_columns, how='left', suffixes=('', '_mean'))
        st.session_state.test[i] = merged_test[i].fillna(merged_test[f"{i}_mean"])
        # 填充val
        merged_val = st.session_state.val.merge(group_means, on=groupby_columns, how='left', suffixes=('', '_mean'))
        st.session_state.val[i] = merged_val[i].fillna(merged_val[f"{i}_mean"])

2. 处理未知分组（可选） #

如果test/val存在train没有的分组，可以用全局均值作为兜底填充：

# 在计算group_means后添加全局均值
global_mean = st.session_state.train[i].mean()
# 先填分组均值，再填全局均值
st.session_state.test[i] = merged_test[i].fillna(merged_test[f"{i}_mean"]).fillna(global_mean)
st.session_state.val[i] = merged_val[i].fillna(merged_val[f"{i}_mean"]).fillna(global_mean)

内容的提问来源于stack exchange，提问作者유정현