局部变量sampled_df赋值前引用问题及DataFrame采样逻辑排查

错误分析 #

• 变量未初始化：sampled_df仅在循环内满足特定条件时才赋值，若所有列都触发continue，或循环未找到符合条件的列，后续引用sampled_df会抛出"赋值前被引用"的错误。
• 循环逻辑错误：
  • 原代码的break位置错误，导致仅检查倒序后的第一列，若该列触发continue，则直接跳出循环，无任何采样操作。
  • 分层采样的条件判断逻辑颠倒：需求中"非分类类型但value_counts长度小于5"应执行分层采样，但原代码中写成了elif len(df[col].value_counts()) >=5，完全违背需求。
• 默认采样缺失：未处理"所有列都不符合分层采样条件"的场景，导致无返回值。

修正后的代码 #

import pandas as pd

tmpdf = pd.DataFrame(data=[[1,2,3,4,5],[2,3,4,5,6],[3,4,5,6,7],[4,5,6,7,8],[5,6,7,8,9]],columns=['a','b','c','d','e'])

def sample_data(df):
    file_size = 2
    sampled_df = None
    
    # 倒序遍历列，寻找符合分层采样条件的列
    for col in df.columns[::-1]:
        col_dtype = df[col].dtype
        value_counts = df[col].value_counts()
        vc_len = len(value_counts)
        # 计算单值占比（出现次数为1的值的总占比）
        single_val_ratio = len(value_counts[value_counts == 1]) / len(df)
        # 计算缺失值占比
        na_ratio = df[col].isna().sum() / len(df)
        
        # 过滤不符合分层采样前置条件的列
        if vc_len == 1 or single_val_ratio > 0.05 or na_ratio > 0.1:
            continue
        
        # 满足分层采样核心条件则执行采样
        if col_dtype in ('object', 'category') or vc_len < 5:
            sampled_df = df.groupby(col, group_keys=False).apply(
                lambda x: x if len(x) < 2 else x.sample(frac=1/file_size, random_state=2022)
            )
            break  # 找到符合条件的列后停止遍历
    
    # 所有列都不符合分层条件时，执行简单随机采样
    if sampled_df is None:
        sampled_df = df.sample(frac=1/file_size, random_state=2022, ignore_index=True)
    
    sampled_df = sampled_df.reset_index(drop=True)
    return sampled_df

# 测试示例
result = sample_data(tmpdf)
print(result)

修正说明 #

• 初始化变量：提前将sampled_df设为None，彻底避免未赋值就引用的问题。
• 修正条件逻辑：
  • 明确分层采样的核心条件：列类型为object/category，或非分类类型但value_counts长度<5。
  • 统一前置过滤条件（单值占比、缺失值占比、value_counts长度），减少重复代码。
• 调整循环逻辑：仅在找到符合条件的列并完成分层采样后才break，当前列不符合则继续遍历下一列。
• 添加默认采样：当所有列都不符合分层采样条件时，自动执行简单随机采样，确保函数总有返回值。
• 代码优化：提前计算value_counts、占比等变量，避免重复计算；将函数名改为更具语义的sample_data，提升可读性。

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