基于行非NA值阈值高效删除DataFrame子集列的方案探究

问题背景 #

我有一个包含1000万行的DataFrame，其详细信息如下：

>>> df.info(show_counts=True)

#   Column  Non-Null Count     Dtype         
---  ------  --------------     -----         
 0   date    10000000 non-null  datetime64[ns]
 1   cust1   6000647 non-null   float64       
 2   cust2   6001585 non-null   float64       
 3   cust3   6000415 non-null   float64       
 4   cust4   9001290 non-null   float64       
 5   cust5   9000402 non-null   float64       
 6   cust6   9000093 non-null   float64       
 7   cust7   8999538 non-null   float64       
 8   cust8   9000211 non-null   float64       
 9   cust9   9000745 non-null   float64       
 10  cust10  9001119 non-null   float64

示例里cust1、cust2、cust3列的NA值占比约40%，其余cust列约10%；date列无缺失值（通用场景下所有列都可能存在任意数量的NA）。

我的需求是：必须保留date列，删除那些非NA值行数低于总行数70%（即700万行）的custXX列，需要一种符合Python惯用写法且高效的实现方式。

基准方案（全表应用dropna） #

我先测试了直接对整个DataFrame使用dropna(axis=1, thresh=thresh)的耗时作为基准：

%timeit df.dropna(axis=1, thresh=thresh)

701 ms ± 12.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

但这个方案有缺陷：它会对所有列（包括date）应用阈值过滤，如果通用场景下date列的非NA值低于阈值，会被误删，不符合我的需求。而dropna的subset参数仅用于行级过滤，无法实现指定列的列级筛选。

我尝试的几种替代方案 #

方案1：拆分DataFrame后合并 #

将DataFrame拆分为仅含date的子表和仅含custXX的子表，对后者过滤后再通过索引合并：

def split_merge(df):
    date_df = df[['date']]
    rest_df = df.drop('date', axis=1)
    cleared = rest_df.dropna(thresh=thresh, axis=1)
    return date_df.merge(cleared, left_index=True, right_index=True)

耗时测试结果：

1.65 s ± 49.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

方案2：统计非NA值后筛选列 #

先筛选出所有custXX列，统计各列非NA值数量，筛选出符合阈值的列后从原表中选取：

def count_select(df):
    nan_cols = df.filter(like='cust').columns
    non_na_counts = df[nan_cols].notna().sum()
    valid_cols = non_na_counts[non_na_counts >= thresh]
    all_cols = pd.concat([pd.Series(0, index=['date']), valid_cols]).index
    return df[all_cols]

耗时测试结果：

1.73 s ± 79.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

方案3：先过滤cust列再选取 #

与方案2逻辑类似，但先对custXX列使用dropna过滤，再提取有效列名从原表中选取：

def select_dropna_select(df):
    nan_cols = df.filter(like='cust')
    cleared = nan_cols.dropna(axis=1, thresh=thresh).columns
    new_cols = ['date', *cleared.values]
    return df[new_cols]

耗时测试结果：

1.54 s ± 14 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

疑问 #

方案3是我尝试的几种里最快的，但仍比基准方案慢一倍以上。请问有没有符合Python惯用写法、且效率接近基准方案的实现方式？

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