Dask and Numba - How to use map partitions efficiently?

Viewed 1462

I'm trying to speed up my code and improve my understanding of Dask and Numba, I did try to use both on some example I created yet there is no improvement and I can't understand why.

I must say that I'm on on a laptop with four cores so the improvement might not be large but it should be there nonetheless.

More precisely on a Windows 10 laptop, using Python 3.7 and having Numba and Dask in a conda environment.

Here is my code :

import numpy as np
import pandas as pd
from numba import jit
import dask.dataframe as dd

data = np.random.randint(-10, 10, (10**8, 3))
df = pd.DataFrame(data=data, columns=["A", "B", "C"], index=None)
df["F"] = np.random.choice(["apple", "banana", "orange", 
                            "pear","grape","lime","citrus","peach"],10**8)

As you can see it's quite a big dataframe memory wise, this was my way to check if Dask represented on improvement or not. On smaller dataframes (< 200MB) it didn't seem to make a difference at all.

ddf = dd.from_pandas(df,npartitions=12)

@jit
def remove_special_char_with_numba(x):
    return x.replace('r','')

This is a toy example where I try to remove a string from a specific column, Numba does speed up the code compared to vanilla Pandas but doesn't support strings so I can't modify replace or use the nopython mode. Now :

%%timeit
remove_special_char_with_numba(df["F"])

output :

58.9 s ± 9.51 s per loop (mean ± std. dev. of 7 runs, 1 loop each)

Next, my understanding of the following is that Dask separated my dataframe into different chunks/partitions and that it will apply the function on each separated chunk independently. With four cores, it should speed up the process as far as I understand it.

%%timeit
ddf["F"].map_partitions(remove_special_char_with_numba).compute()

output :

45.9 s ± 10.5 s per loop (mean ± std. dev. of 7 runs, 1 loop each)

Now I don't want to be greedy but shouldn't the improvement be larger than that? Am I doing something wrong ?

Thanks

1 Answers

This result should not surprise you too much. You are, apparently, running on the default threaded scheduler.

That means that every string operation must obtain the single python GIL in order to happen, and that is true whether it is in dask-controled worker threads. That remains true for the numba-jit version of the operation, because you cannot run this function in no-python mode. If it were in no-python mode, it would release the GIL, and full string support will be coming to numba.

You might be able to get a better speedup using the distributed scheduler with a number of processes, although then you suffer the costs of sending data between processes, so it matters how you generate the data, and the fact that you bring all of the results into the main session when you compute().

Related