日々精進

新しく学んだことを書き留めていきます

daskで何故か使用メモリ量が大幅に増えるケースがある

手元のデータセットで以下のコードを実行すると使用メモリ量が2GB程度増えた。

    def _get_difference(self, num_feat_names: List[str], num_feats: DataFrame) -> DataFrame:
        def diff(df: DataFrame) -> DataFrame:
            return df.diff(1)
        npartitions: int = min(os.cpu_count() * 3, int(len(num_feats) / 100))
        ddf: dd.DataFrame = dd.from_pandas(num_feats, npartitions=npartitions).set_index("customer_ID")
        ddf = ddf[num_feat_names].map_overlap(diff, 2, 0)
        ddf = ddf.rename(columns={col: col + "_diff_1" for col in num_feat_names})
        ddf = ddf.groupby("customer_ID").last()
        return ddf.reset_index().compute()

上記コードの、

ddf = ddf[num_feat_names].map_overlap(diff, 2, 0)

ddf[num_feat_names] = ddf[num_feat_names].map_overlap(diff, 2, 0)

に変更すると使用メモリ量の増加が24GB程度になった(!)なぜだ・・

compute()した後の行数・列数は同じだし、ddfオブジェクトを削除してGCを実行しても使用メモリ量はほぼ変わらず。

daskが内部で持っているデータを解放できていないように見える。

こんなケースがあるとdask使いづらいな。。