dask DataFrame.assign blows up dask graph

Question

So I have an issue with dask DataFrame.append. I generate a lot of derivative features from main data and append them to the main dataframe. After that the dask graph for any set of columns is blown up. Here is small example:

%pylab inline
import numpy as np
import pandas as pd
import dask.dataframe as dd
from dask.dot import dot_graph

df=pd.DataFrame({'x%s'%i:np.random.rand(20) for i in range(5)})
ddf = dd.from_pandas(df, npartitions=2)

dot_graph(ddf['x0'].dask)

here is the dask graph as expected

g=ddf.assign(y=ddf['x0']+ddf['x1'])
dot_graph(g['x0'].dask)

here the graph for same column is exploded with irrelevant computation

Imagine i have lots of lots of spawned columns. So computation graph for any particular column includes irrelevant computations for all the other columns. I.e. in my case I have len(ddf['someColumn'].dask)>100000. So that becomes unusable quickly.

So my question is can this issue be resolved? Are there any existing means to do this? If not - what direction should i look to implement this?

Thanks!

Answer 1

Rather than continuously assigning new columns to the dask dataframe, you might want to build several dask series and then concat them all together at the end

So instead of doing this:

df['x'] = df.w + 1
df['y'] = df.x * 10
df['z'] = df.y ** 2

Do this

x = df.w + 1
y = x + 10
z = y * 2
df = df.assign(x=x, y=y, z=z)

Or this:

dd.concat([df, x, y, z], axis=1)

This may still result in the same number of tasks in your graph however, but will probably result in fewer memory copies.

Alternatively, if all of your transformations are row-wise then you can construct a pandas function and map that across all partitions

def f(part):
    part = part.copy()
    part['x'] = part.w + 1
    part['y'] = part.x * 10
    part['z'] = part.y ** 2
    return part

df = df.map_partitions(f)

Also, while a million-node task graph is less than ideal, it should also be OK. I've seen larger graphs run comfortably.