Dask分区远小于Worker内存仍出现内存溢出问题求助

在项目中使用Dask时，已遵循以下最佳实践：

• 用map_partitions封装原生Pandas代码
• 为合并/分组列创建索引
• 设置分区大小小于Worker内存
• 避免大任务图（如建索引后及时保存）

但用随机数据测试时，即使总数据量仅2.4GB（远小于本地16GB内存），单分区约23MB（远小于Worker分配的4GB内存），仍出现Worker内存超限重启、本地内存耗尽的情况。已增大索引列取值范围，排除合并时重复值导致数据膨胀的可能，求问原因及解决办法。

import pandas as pd
from dask.distributed import Client, LocalCluster
import dask.dataframe as dd 
import numpy as np


def create_files():
    size = 1_000_000
    for i in range(100):
        df = pd.DataFrame({"col1": np.random.randint(90_000, 100_000, size), "col2": np.random.randint(101, 20_000, size), "col3": np.random.uniform(0, 10_000, size)})
        # Select appropriate partitions
        ddf = dd.from_pandas(df, npartitions=1)
        ddf.to_parquet(f"test", ignore_divisions=True, engine="fastparquet", overwrite= i==0, append= i>0)

    print("Created first file")
    for i in range(10):
        df = pd.DataFrame({"col1": np.random.randint(90_000, 100_000, size), "col4": np.random.uniform(0, 10_000, size)})
        # Select appropriate partitions
        ddf = dd.from_pandas(df, npartitions=1)
        ddf.to_parquet(f"test2", ignore_divisions=True, engine="fastparquet", overwrite= i== 0, append= i>0)
    print("Created second file")
    print("-------------------------------------------------------------")
def index_reparttion():
    print("About to repartition")
    ddf = dd.read_parquet("test")
    ddf = ddf.repartition(npartitions=100)
    ddf = ddf.set_index("col1")
    ddf = _rebalance_ddf(ddf)
    print("save parquet")
    ddf.to_parquet("test")
    

    ddf = dd.read_parquet("test2")
    ddf = ddf.repartition(npartitions=50)
    ddf = ddf.set_index("col1")
    ddf = _rebalance_ddf(ddf)
    print("save parquet 2")
    ddf.to_parquet("test2")
        
# https://stackoverflow.com/questions/52642966/repartition-dask-dataframe-to-get-even-partitions
def _rebalance_ddf(ddf):
    """Repartition dask dataframe to ensure that partitions are roughly equal size.

    Assumes `ddf.index` is already sorted.
    """
    if not ddf.known_divisions:  # e.g. for read_parquet(..., infer_divisions=False)
        ddf = ddf.reset_index().set_index(ddf.index.name, sorted=True)
    index_counts = ddf.map_partitions(lambda _df: _df.index.value_counts().sort_index()).compute()
    index = np.repeat(index_counts.index, index_counts.values)
    divisions, _ = dd.io.io.sorted_division_locations(index, npartitions=ddf.npartitions)
    return ddf.repartition(divisions=divisions)


def main():
    ddf = dd.read_parquet("test")
    print(ddf.compute())
    print(ddf.memory_usage_per_partition(index=True, deep=False).compute())
    print(ddf.memory_usage(deep=True).sum().compute())
    
    ddf2 = dd.read_parquet("test2")
    print(ddf2.memory_usage_per_partition(index=True, deep=False).compute())
    print(ddf2.memory_usage(deep=True).sum().compute())
    def mapped_fun(data):
        for lag in range(4):
            data[f"col_{lag}"] = data.groupby("col1")["col3"].transform(lambda x: x.shift(lag)).apply(lambda x: np.log(x))
        return data

    ddf = ddf.map_partitions(mapped_fun)
    
    ddf = ddf.merge(ddf2, on=['col1'], how="left")

    ddf.to_parquet("result", engine="fastparquet")


if __name__ == "__main__":
    cluster = LocalCluster(
                    n_workers=4, 
                    threads_per_worker=2,
                    memory_limit='auto'
                    )
    client = Client(cluster)
    create_files()
    index_reparttion()
    main()

1. 不必要的全量compute()操作：
   • main函数中的ddf.compute()会把1亿行的test数据集全量加载到本地内存，Parquet是压缩格式，解压后内存占用会翻倍甚至更多，直接耗尽本地内存。
   • _rebalance_ddf中的index_counts.compute()和np.repeat会生成1亿元素的数组，仅这个数组就占用约800MB内存，后续处理还会进一步消耗资源。
2. 分区内计算的内存膨胀：
   • mapped_fun中按col1分组移位后新增4列，每个分区的内存占用会显著增加；同时apply(np.log)未处理shift产生的NaN值，可能触发额外内存开销。
3. Cluster内存配置不合理：
   • memory_limit='auto'会让Dask高估Worker可用内存，4个Worker共享16GB内存时，容易出现内存竞争导致超限重启。

1. 移除全量compute()，改用采样查看：
   • 删除main中的print(ddf.compute())，如需验证数据，改用ddf.head(10).compute()或ddf.sample(100).compute()，仅加载少量样本。
   • 替换_rebalance_ddf的手动分区逻辑，用Dask内置的repartition(partition_size="20MB")自动平衡分区，避免生成大数组：

     def index_reparttion():
         print("About to repartition")
         ddf = dd.read_parquet("test")
         ddf = ddf.set_index("col1").repartition(partition_size="20MB")
         ddf.to_parquet("test")
         
         ddf = dd.read_parquet("test2")
         ddf = ddf.set_index("col1").repartition(partition_size="20MB")
         ddf.to_parquet("test2")
     

2. 优化分区内计算逻辑：
   • 提前处理shift产生的NaN值，合并分组操作减少中间变量：

     def mapped_fun(data):
         grouped = data.groupby("col1")["col3"]
         for lag in range(4):
             # 先填充空值再计算log，避免异常
             data[f"col_{lag}"] = np.log(grouped.shift(lag).fillna(1))
         return data
     

3. 显式配置Worker内存：
   • 给每个Worker设置合理的内存上限，避免内存竞争：

     cluster = LocalCluster(
         n_workers=4, 
         threads_per_worker=2,
         memory_limit="3GB"
     )
     

4. 减少重复I/O操作：
   • 在create_files中直接生成带分区的Parquet，无需后续重新分区后覆盖写入，降低I/O和内存开销。

内容的提问来源于stack exchange，提问作者J.Ewa