Use atop fusion in dask dataframe

[mrocklin]

The high level atop fusion layers can be used in dask dataframe as well as dask array. Today this makes it cheaper to perform task fusion on dataframes with many partitions. In the future this should make it easier to build more sophisticated high level optimizations.

This currently builds on #4092

This currently breaks tests. There is a fair amount that we can clean up in both dask.dataframe's broadcast operations (map_partitions, elemwise) and in the dask.array.atop code.

• Tests added / passed
• Passes flake8 dask

[mrocklin]

OK, most dataframe tests pass here now

[mrocklin]

OK, this is ready for review. There is some future work though:

• Regression: I've turned off dataframe-parquet column projection for now. (The previous implementation was very brittle). I think that we'll want to redo it anyway in short order.
• I think that we should rename atop to blockwise, or something else (see Better name for atop? #4035)
• There are issues around fusing diamond-like atop graphs. We might need to repeat some of the work of @jcrist and @eriknw at a higher level.

[mrocklin]

@jcrist @TomAugspurger if either of you have time it would be helpful if you could take a look at this

[TomAugspurger]

Will probably have time next week.

…

On Wed, Dec 12, 2018 at 5:13 PM Matthew Rocklin ***@***.***> wrote: @jcrist <https://github.com/jcrist> @TomAugspurger <https://github.com/TomAugspurger> if either of you have time it would be helpful if you could take a look at this — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#4229 (comment)>, or mute the thread <https://github.com/notifications/unsubscribe-auth/ABQHIh_ilQHPW7ecby3zsAWyGHgDRl91ks5u4Y2agaJpZM4Yrofg> .

[mrocklin]

This is future work. We need better fusion on the high level graph level. cc @jcrist ?

[jcrist]

Sure, that'd be fun to work on. Is this blocking for this PR to merge, or just future work?

[mrocklin]

Not blocking. Things are better with this PR (df + 1 + 2 + 3 + 4 happens in one task) but could be much better still with better fusion. In particular I've seen the use case in this failing test happen frequently. It seems that common use of Pandas includes dozens of lines of modifying columns in place, all of which generate diamond-like graphs.

[mrocklin]

The PR itself could also use review though.

[TomAugspurger]

This PR looks good, but I haven't really kept up with the new high-level graph stuff.

Only request would be a docstring for the new broadcast method. And could you explain the reasoning behind that name? IIUC, adds a task to each partition / block of the input args? I worry a bit about broadcast because it clashes with NumPy's concept (and broadcasting data in distributed), but I haven't come up with a better name.

[mrocklin]

A fun motivating stack overflow question: https://stackoverflow.com/questions/53844188/how-do-i-use-dask-to-efficiently-calculate-many-simple-statistics/53869772#53869772

[mrocklin]

Only request would be a docstring for the new broadcast method. And could you explain the reasoning behind that name?

Renamed to partitionwise_graph and added a docstring. I also moved it into core.py

[bluecoconut]

Just did a test locally and this branch had a dramatic impact (~2x improvement) in performance for the code below, and even has a lower memory footprint for the larger stat tests.

import dask
import dask.datasets
import numpy as np
import time
from distributed import Client

client = Client()
client

df = dask.datasets.timeseries()
df = df.repartition(npartitions=300)
df = client.persist(df)

def random_indexer(df):
    indexer = ~df.index.isnull()
    for i in range(np.random.randint(15)+1):
        col = np.random.choice(['x','y'])
        value = np.random.uniform(-1,1)
        op = np.random.choice([lambda x, y: x < y, lambda x, y: x > y])
        indexer = np.logical_and(indexer, op(df[col], value))
    return indexer

def random_statistic(indexer, df):
    col = np.random.choice(['x', 'y', 'name'])
    if col == 'name':
        op = np.random.choice([lambda x: x.unique().size, np.min, np.max])
    else:
        op = np.random.choice([lambda x: x.unique().size, np.min, np.max, np.sum, np.mean])
    return op(df[col][indexer])

np.random.seed(137)
stats = []
for i in range(10):
    ind = random_indexer(df)
    for k in range(20):
        stats.append(random_statistic(ind, df))

st = time.time()
stat_computed = client.compute(stats)
ft = time.time()
print(ft-st)

st = time.time()
stat_results = client.gather(stat_computed)
ft = time.time()
print(ft-st)

For calculating 200 statistics

(10 unique filtering, 20 different statistics from that filtered subset)

npartitions	graph create (master)	graph create (high-level)	execution (master)	execution (high-level)
100	1.7s	2.5s	10.2s	5.6s
300	5.4s	8.1s	31s	15.3s
600	11.6s	16s	59 s	29.35s

For calculating 2000 statistics

I also tried to increase the scale of stats (for i in range(100), 2000 stats total) and I see the same garbage collection warnings in both branches, and memory still seems to grow a lot more than I anticipated.

However, on this branch, the calculation with 2000 stats actually completes! Also, the Bytes stored on the dashboard seems to be accurate (though quite large, from 1.5 GB after persist(df) to 7 GB stored after calculating stats. (will need to dig into why there is so much data left in the cluster next).

The biggest difference from current 1.0.0 master is that master actually doesn't complete this same calculation on my machine. It has the memory issue (?) that even prevents this from completing, heading into swap and eventually killing workers.

[mrocklin]

Thank you for trying things out @bluecoconut and for the benchmark. It's always nice to see things work well on unseen problems :)

I look forward to finding out why graph construction was more expensive (I would have expected it to decrease).

I suspect that you'll see additional boosts as we get a bit better at fusing things at the high level (see the comments on the xfailed test).

[mrocklin]

Merging this tomorrow if there are no further comments.