Filtering chunks of a dataframe in parallel using Dask

user_2690677 0
Viewed 435

I have a use case where I need to filter pieces of a data frame belonging to some factor group.

For instance, lets say the data frame looks about like this:

index    col1    col2
0        A       5
1        A       10
2        B       5
3        B       8
4        C       0
... and so on

What I need to do is take a chunk of the data where df[df['col1'] == 'A'], then apply another function that tests the rows on some other criteria and further filters down the data.

So my loop currently looks like this:

df_chunks = [df[df['col1'] == x] for x in col1_unique_values]
outputs = []
for df_chunk in df_chunks:
    output = do_my_custom_filter(df_chunk)
    outputs.append(output)

Is there a way to replace the last for loop with Dask? I have tried the multiprocessing route and found it to be difficult to debug.

Any design patterns to that end would be very helpful!

dask parallel-processing python
1 Answers

It looks like your computation might be embarrassingly parallel. If so, you might write a single function that consumes and produces Pandas dataframes, and then just call df.map_partitions(your_function)

Alternatively, if you want the chunks to be coalesced together then you might want to look at groupby-apply, like the following:

df.groupby("col1").apply(my_custom_function)
user_6166160
Related