Dask DataFrame筛选子集后转Pandas出现内存溢出问题求助

核心原因 #

• 重复执行计算链路，无中间结果缓存
  Dask为懒执行框架，你未对empty_b = df[df['b'].isna()]的中间结果做持久化（调用persist()），后续两次调用value_counts时，每次都会重新执行「全量Parquet读取->b列空值筛选」的完整链路，相当于数据被加载、筛选了两次，直接翻倍了IO和内存开销。
  而第一次全量统计a列时，Dask对Parquet的投影下推优化生效，仅需读取a列即可完成统计，无额外计算开销。
• 投影下推优化失效，加载数据量大幅提升
  第一次全量统计a列时，Dask可直接向Parquet下推投影规则，仅加载a列这1列数据，160M条单列式数据内存占用仅1.5GB左右。但当逻辑变为「筛选b列空值->统计a列」时，计算链路需要同时加载a、b两列才能执行，若未在读取Parquet时显式指定列，甚至会加载全部240列数据，加载的数据量较第一次提升几十倍，直接拉高内存占用。
• 单分区+重复统计叠加shuffle压力
  两次统计的结果最终都汇聚为1个分区，若b列为空的子集占比不低，全量数据汇聚到单分区做频次统计时，单分区内存负载本身就很高，叠加两次value_counts触发的两次全量shuffle汇聚，内存占用直接超过集群阈值触发OOM。

优化方案 #

1. 读取Parquet时显式指定需要的列，触发投影下推，仅加载必要的a、b两列：

import dask.dataframe as dd
# 显式指定columns避免加载全量240列
df = dd.read_parquet(file, columns=['a', 'b'])

2. 筛选完空值子集后先持久化中间结果，避免重复执行筛选链路：

empty_b = df[df['b'].isna()].persist()

3. 仅调用一次value_counts生成计数，再基于计数计算占比，避免触发两次统计链路：

count = empty_b.a.value_counts()
percent = count / count.sum()
a_count = dd.concat([count, percent], axis=1, keys=['counts', '%'])
res = a_count.compute()

4. 若a列基数极高，可在统计前对empty_b做重分区，降低单分区的内存压力。

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