dask bag `take` ignores shared dependencies

Normally, take seems to intelligently “propagate” its way down, so that a take(1) doesn’t evaluate the entire task graph before . This is super useful for exploring large data sets.

Unfortunately, this doesn’t seem to work for “diamond” / shared dependencies (like those created with zip), where the task graph looks like:

     --------
     | Take |
     --------
        /\
      /    \
    /        \
--------  --------
| Task |  | Task |
--------  --------
    \        /
      \    /
        \/
     --------
     | Base |
     --------

In this diamond dependency, even a trivial take(1) seems to evaluate the entire task graph, even though it should be easy to for take(1) to percolate down through the zip.

I saw this on Python 3.5 and dask 0.10.2 (both from conda-forge).

[mrocklin]

Yup, so here’s what’s going on.

Normally dask.bag partitions are iterators, they stream data if possible. This means that it’s even nicer than normal dask collections like dask.dataframe because we only need to compute a tiny bit of a partition to get something like take 1. Because of this streaming-within-partitions ability dask.bag.read_text feels free to read each file as one partition by default.

This is great except when you tee off. When this happens we no longer trust iterators (because each part of the tee would separately consume the same iterator) so we reify everything into a list. Now we’re still good about only computing the right partitions but we’re bad about being lazy within each partition.

Unfortunately now our partition-per-file choice becomes bad, we’re reading the entire file. You can avoid this choice by explicitly sending in a chunksize=10000000 or something, which will break the file up into 10MB blocks. Then take(1) will only process 10000000 bytes instead of the entire file.

Does this help?

[jakirkham]

It seems like it should still be possible (maybe even easier today) to optimize this case. Is there any interest in taking this on?