Python函数中处理DataFrame NaN值：SQL查询结果异常修复

问题场景 #

我的DataFrame中localdf["where_condition"]列包含有效值与NaN：

0    FirstName ='sonali'
1             Gender='F'
2                   NaN

编写的统计函数如下：

def filter_record_count_check(localdf):
    try:
        sql = 'select count(*) as CNT from ' + localdf["schema_table"] + ' where ' + localdf["where_condition"]
        sql = sql.to_dict()
        print(sql)
        df_list=[]
        for i in sql.values():
            df = pd.read_sql_query(i, db_connection)
            df_list.append(df.CNT[0])
            print("df_list")
            print(df_list)
        return df_list

生成的SQL字典包含NaN项：

{0: "select count(*) as CNT from testdb.DimCurrency where FirstName ='sonali'", 1: "select count(*) as CNT from testdb.banking_fraud where Gender='F'", 2: nan}

执行后得到的df_list为[2, 208]，但将该结果赋值到原DataFrame的filter_column_cnt列时，整列都变成了NaN，无法得到期望的对应行匹配结果：

当前错误输出 #

table_name column_name  ... record_count  filter_column_cnt
0    DimCurrency    PersonID  ...            6                NaN
1  banking_fraud         NaN  ...         1000                NaN
2  banking_fraud         NaN  ...         1000                NaN

期望输出 #

table_name column_name  ... record_count  filter_column_cnt
0    DimCurrency    PersonID  ...            6                2
1  banking_fraud         NaN  ...         1000                208
2  banking_fraud         NaN  ...         1000                NaN

问题原因 #

原函数存在两个核心问题：

1. 当where_condition为NaN时，拼接出的SQL为NaN，执行pd.read_sql_query会抛出异常，导致该行的统计值未被添加到df_list，最终返回的列表长度（2）与原DataFrame行数（3）不匹配，赋值时触发全列NaN。
2. 直接遍历SQL字典的values()，无法保证与原DataFrame的行顺序完全对应（虽本例中顺序一致，但存在潜在风险）。

修正后的代码 #

遍历原DataFrame的每一行，针对where_condition的有效值/NaN分别处理，确保返回的列表长度与原DataFrame完全匹配：

import pandas as pd

def filter_record_count_check(localdf):
    df_list = []
    for _, row in localdf.iterrows():
        # 判断where_condition是否为NaN
        if pd.isna(row["where_condition"]):
            df_list.append(pd.NA)
            continue
        
        # 拼接有效SQL语句
        sql = f'select count(*) as CNT from {row["schema_table"]} where {row["where_condition"]}'
        try:
            # 执行SQL查询并提取统计值
            df = pd.read_sql_query(sql, db_connection)
            df_list.append(df.CNT.iloc[0])
        except Exception as e:
            # 捕获SQL执行异常，可根据需求调整处理逻辑
            print(f"SQL执行失败: {sql}, 错误信息: {str(e)}")
            df_list.append(pd.NA)
    return df_list

效果验证 #

调用修正后的函数，返回的df_list为[2, 208, pd.NA]，将其赋值到原DataFrame的filter_column_cnt列后，即可得到与期望一致的输出：有效条件行显示统计值，NaN条件行保留NaN。

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