Repository: dask/dask

Issue:
Best way to iterate over groups
hello,
I'm trying to translate a custom join logic from pandas over to dask, I would like to some ideas about what would be the most optimal way of doing it. 

The input to the function are left and right tables that have both timestamp columns and a key. For each key, the function matches up a row in the left table with the latest row in the right table that is just before it in time. 

This is how I had implemented it in pandas:

``` python
def __pandas_join__(left, right, key, lts, rts):
   """
  :param lts,rts : name of timestamp column in left and right tables. 
                       Assumed to be int64 types
   """
       lgroups = left.groupby(key)[lts];
       rgroups = right.groupby(key)[rts];

        both_groups = [ ];
        #iterate over all keys in the left table
        for k, ldf in lgroups:
        try:
                    rdf = rgroups.get_group(k)
        except KeyError as e: # the key does not exist in the right table
                    rdf = Series( nan, index=ldf.index )
        both_groups.append( ( k, ldf, rdf ) )

        # __leftjoin__mapper__ is a numba.jit(nopython=True, nogil=True) function that
        # for every row in the left table calculates the correct index in the right table. 
        # It returns a numpy array the length of the left table, giving the correct indices in
       # the right table
        left['__lrix_map__'] = concat( [
                Series(
                    __leftjoin_mapper__(
                        ldf.as_matrix().astype(np.float), # nopython mode in numba.jit
                        rdf.as_matrix().astype(np.float), # requires typed numpy arrays
                        rdf.index.values,
                        tolerance, strict
                    ),
                    index=ldf.index
                )
                for k, ldf, rdf in both_groups
        ] )

        merge_df = merge( left, right, left_on='__lrix_map__', right_index=True, suffixes=suffixes );

       return merge_df
```

This code takes about a minute to join two tables with 10mn rows. There is a lot of opportunity of embarrassing parallelism and I wanted to take advantage of dask's scheduling to scale this to datasets 100+mn rows. However, I'm not able to figure out an efficient way to iterate over groups.

I've tried something like this 

``` python
from multiprocessing.pool import Pool, ThreadPool
(mapper, mapper_name) = (ThreadPool().starmap, 'ThreadPool.starmap')

keys = unique(left[key].map_partitions(unique).compute())

@jit(nogil=True)
def lrmapper(k, lg, rg, rix,  tol, strict):    
    # sort the left timestamps and presever the map back to the 
    # original mutation
    lsort_ix = argsort(lg)
    replace_ix = argsort(lsort_ix)

    #sort the right timestamps
    rsort_ix = argsort(rg)

    lrix_map = pcjn.__leftjoin_mapper__(lg[lsort_ix], rg[rsort_ix], rix[rsort_ix], tol, strict)    
    return k, lrix_map[replace_ix]


args = [];
lgroups = left.groupby(key)
rgroups = right.groupby(key)

    args = ( 
        (
            k, 
            lgroups.get_group(k)[lts].compute().as_matrix().astype(float64), 
            rgroups.get_group(k)[rts].compute().as_matrix().astype(float64),
            rgroups.get_group(k)[rts].index.compute().values
        )
        for k in keys
    )

# parallel execution: my patience runs out before it finishes.
# import profile 
# %prun mapper(lrmapper, args)

# serial execution --- takes about 5 secs per key 
for a in args:
    lrmapper(*a)
```

This serial version of the code takes about 10mins for a dataset with 3000 keys and 200 rows per key (ie ~6e5 rows), while the parallel version literally takes for ever. 
Profiling the multithreaded version on just 10 keys, it appears that the most of the time is spent waiting for the lock.

```
 ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     9742   44.061    0.005   44.061    0.005 {method 'acquire' of '_thread.lock' objects}
    13960    0.161    0.000    0.381    0.000 core.py:179(get_dependencies)
      420    0.126    0.000    0.126    0.000 {built-in method stat}
    99020    0.067    0.000    0.078    0.000 core.py:27(istask)
    40460    0.059    0.000    0.110    0.000 core.py:147(_deps)
    98080    0.058    0.000    0.063    0.000 {built-in method isinstance}
     2260    0.036    0.000    0.079    0.000 async.py:299(finish_task)
     2260    0.035    0.000    0.255    0.000 async.py:445(fire_task)
       30    0.031    0.001   44.225    1.474 async.py:374(get_async)
        1    0.030    0.030    0.030    0.030 {built-in method listdir}
     2260    0.029    0.000   43.584    0.019 queue.py:150(get)
```

Any suggestions / ideas on how to do this more efficiently would be greatly appreciated. 

Thanks !



Discussion:
It'll take me a little while to sift through this code.  I may not have time to deal with this today.


Resolving patch:
diff --git a/dask/array/optimization.py b/dask/array/optimization.py
--- a/dask/array/optimization.py
+++ b/dask/array/optimization.py
@@ -22,7 +22,7 @@ def optimize(dsk, keys, **kwargs):
                              set([getarray, getitem, np.transpose]))
     dsk2 = cull(dsk, list(flatten(keys)))
     dsk3 = remove_full_slices(dsk2)
-    dsk4 = fuse(dsk3)
+    dsk4 = fuse(dsk3, keys)
     dsk5 = valmap(rewrite_rules.rewrite, dsk4)
     dsk6 = inline_functions(dsk5, fast_functions=fast_functions)
     return dsk6
diff --git a/dask/bag/core.py b/dask/bag/core.py
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -100,7 +100,7 @@ def inline_singleton_lists(dsk):
 def optimize(dsk, keys, **kwargs):
     """ Optimize a dask from a dask.bag """
     dsk2 = cull(dsk, keys)
-    dsk3 = fuse(dsk2)
+    dsk3 = fuse(dsk2, keys)
     dsk4 = inline_singleton_lists(dsk3)
     dsk5 = lazify(dsk4)
     return dsk5


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -1744,3 +1744,11 @@ def test_chunks_error():
 def test_array_compute_forward_kwargs():
     x = da.arange(10, chunks=2).sum()
     x.compute(bogus_keyword=10)
+
+
+def test_dont_fuse_outputs():
+    dsk = {('x', 0): np.array([1, 2]),
+           ('x', 1): (inc, ('x', 0))}
+
+    a = da.Array(dsk, 'x', chunks=(2,), shape=(4,))
+    eq(a, np.array([1, 2, 2, 3], dtype=a.dtype))