如何将Dask布尔标量转为布尔表达式并解决内存溢出问题

问题背景 #

处理超大规模数据集时，通过Dask Series比较得到布尔值的过程中出现内存溢出，且预期的布尔值实际为Dask Series类型。核心代码片段如下：

scores_for_test_data[f"Healthy Sample Score {exclude_name}"] = np.power(
    (pairwise_test_excluded_sample - healthy_sample_mean_ratio), 2)

scores_for_test_data[f"Lung Sample Score {exclude_name}"] = np.power(
    (pairwise_test_excluded_sample - lung_sample_mean_ratio), 2)

scores_for_test_data[f"Sum Healthy Sample Score {exclude_name}"] = scores_for_test_data[
    f"Healthy Sample Score {exclude_name}"].sum()

scores_for_test_data[f"Sum Lung Sample Score {exclude_name}"] = scores_for_test_data[
    f"Lung Sample Score {exclude_name}"].sum()

res = scores_for_test_data[f"Sum Healthy Sample Score {exclude_name}"][0] < \
      scores_for_test_data[f"Sum Lung Sample Score {exclude_name}"][0]

if res.any().compute():
    check += 1

所有变量均为Dask对象，求和后的Series所有行值完全重复（示例如下）：

0     2375.861075
1     2375.861075
2     2375.861075
3     2375.861075
4     2375.861075
5     2375.861075

执行时触发内存相关报错，示例日志：

2023-10-22 14:12:56,258 - distributed.worker.memory - WARNING - Unmanaged memory use is high. This may indicate a memory leak or the memory may not be released to the OS; see https://distributed.dask.org/en/latest/worker-memory.html#memory-not-released-back-to-the-os for more information. -- Unmanaged memory: 5.70 GiB -- Worker memory limit: 7.98 GiB
2023-10-22 14:12:56,716 - distributed.worker.memory - WARNING - Worker is at 80% memory usage. Pausing worker.  Process memory: 6.40 GiB -- Worker memory limit: 7.98 GiB
...（省略后续重复日志）

解决方案 #

1. 移除冗余的全量重复列计算 #

原代码中对分数列求和后，将结果广播为全量重复的新列，这会完全浪费内存。直接计算分数列的总和得到Dask标量即可，无需存入DataFrame：

# 计算分数列的总和（直接得到Dask scalar，而非全量重复列）
sum_healthy = scores_for_test_data[f"Healthy Sample Score {exclude_name}"].sum()
sum_lung = scores_for_test_data[f"Lung Sample Score {exclude_name}"].sum()

2. 直接比较标量并获取单个布尔值 #

Dask标量的比较结果仍是Dask标量，直接compute()即可得到Python原生布尔值，无需处理Series：

# 比较两个标量，compute得到原生布尔值
res = (sum_healthy < sum_lung).compute()

# 直接使用布尔值判断
if res:
    check += 1

这样避免了加载全量重复的求和列，仅计算并获取两个总和值，内存占用极低。

3. 额外内存优化建议 #

• 避免在DataFrame中存储中间计算结果：如果Healthy Sample Score和Lung Sample Score仅用于求和，可直接链式计算，无需存入DataFrame：

  sum_healthy = np.power(pairwise_test_excluded_sample - healthy_sample_mean_ratio, 2).sum()
  sum_lung = np.power(pairwise_test_excluded_sample - lung_sample_mean_ratio, 2).sum()
  

• 调整Dask worker内存参数：如果仍有内存压力，可降低worker的内存限制阈值，或增加worker数量分散计算负载。

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