Higher than expected memory usage for repartition

[syagev]

I find that repartitioning a dataframe by splitting is causing higher than expected memory usage. I've encountered this in my project where I do the "not so standard" thing of storing entire numpy arrays (and other objects) inside cells of the dataframe.

The problem is exemplified in this notebook: https://gist.github.com/syagev/4de6f6c1cb25d2b4e9e6a6d846d12793

When processing a large dataframe (where each partition just fits into worker memory) everything is well. When splitting (repartitioning) this dataframe workers fail due to memory usage. However, when converting the objects inside the dataframe to bytes before repartitioning then all is well again. This caused me to think that perhaps this is serialization issue caused by the fact there are objects inside cells in the dataframe.

Is this scenario supported? Should dataframes contain only numeric/string/bytes items?

[mrocklin]

I believe that Dask serializes Pandas dataframes by just pickling them. Numpy arrays serialize decently well with pickle, about the same size as a bytes object would serialize, so I don't know what would case your issue unfortunately.

[syagev]

Thanks for the quick response @mrocklin! As a workaround, can I somehow control the assignment of tasks to avoid serialization altogether? In the scenario where we repartition from N partitions to NxP partitions it makes more sense anyhow to put the split + get_item + reset_index tasks that relate to a specific partition on the worker that holds that partition in the first place. Am I missing something?

[mrocklin]

Yes. In that situation I wouldn't expect any communication of the intermediate data to occur. Does it? It would be visible on the dashboard as red bars in the task stream plot.

…

On Tue, Dec 31, 2019 at 1:37 AM Stav ***@***.***> wrote: Thanks for the quick response @mrocklin <https://github.com/mrocklin>! As a workaround, can I somehow control the assignment of tasks to avoid serialization altogether? In the scenario where we repartition from N partitions to NxP partitions it makes more sense anyhow to put the split + get_item + reset_index tasks that relate to a specific partition on the worker that holds that partition in the first place. Am I missing something? — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#3347?email_source=notifications&email_token=AACKZTCH7DWYXKOZWJEPGT3Q3MHFNA5CNFSM4KBQLIGKYY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOEH37FBQ#issuecomment-569897606>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AACKZTDWJK5VYI75TIRW4CLQ3MHFNANCNFSM4KBQLIGA> .

[syagev]

Yes it does. Indeed, huge red bars on the dashboard.

The split + get_item combo is extremely communication intensive. My partitions are ~1.5GB in size (because I have a mapping function that creates output about 20x the size of the input). Before writing to files I wanted to repartition back into managable partition size (of about 100 MB). I cannot start with smaller partitions as input since my mapping function requires a minimum amount of rows to be meaningful. Since memory usage goes high when partitions are serialized and communicated this causes memory errors.

There are 2 solutions IMO:

1. Somehow force the repartitioning (even-split) tasks to be executed on the worker that has the original partition (maybe using the new Actor concept?)
2. If after the split the next step is serialization to disk (i.e. to_parquet) then enable splitting internally during that stage.

I actually implemented (2) on my private fork - added a split_factor parameter to to_parquet that writes N parquet files for each partition (and returns metadata for N row-groups). This seems to work for me at this point but (1) seems the cleaner approach..

[mrocklin]

Somehow force the repartitioning (even-split) tasks to be executed on the worker that has the original partition (maybe using the new Actor concept?)

Given the size of the data this should be happening automatically (actors are definitely overkill here)

My next guess is that Dask is somehow underestimating the size of these partitions. Can I ask you to get one of these partitions

part = df.partitions[0].compute()

And then run the dask.sizeof function on it to see how Dask estimates its size?

dask.sizeof.sizeof(part)

Is this 1.5GB, or is it something much smaller?

If it's smaller, then perhaps Dask is underestimating size, and so moving things around when it shouldn't.

[syagev]

Yes, much smaller...
it's 28137698 (thats' for ~1Gb partition)

[mrocklin]

Here is the code that we use to estimate the size of a Pandas dataframe:

https://github.com/dask/dask/blob/c69d10f97cadd06f147817f0084bddeadb8195e6/dask/sizeof.py#L80-L97

For object dtype columns (which presumably you have) we randomly select twenty elements, compute the size of those elements, and extrapolate from there. Would you be interested in running through this code on your dataframe and seeing where this estimation process is breaking down?

[syagev]

Will do. Thanks

[jakirkham]

FWIW we added pickle protocol 5 support in Distributed 2.21.0, which supports out-of-band buffers. This cuts down on the memory usage during pickling quite a bit. The result being spilling and data transmission should be more efficient. So it may be worth exploring. Would be curious to know whether the issue still remains.