Repository: dask/dask

Issue:
Synchronize keys between two dasks
Given two dasks

``` python
d = {'a': 1, 'b': (add, 'a', 10), 'c': (mul, 'b', 5)}
e = {'x': 1, 'y': (add, 'x', 10), 'z': (mul, 'y', 2)}
```

We may want to merge them in a way that is blind to renamings of variables.

``` python
>>> merge_sync(d, e)
{'a': 1, 'b': (add, 'a', 10), 'c': (mul, 'b', 5), 'z': (mul, 'b', 2)}
```

I think that @eriknw may have brought up this idea before.  I now have a driving use case for it and find it very valuable.

I'm not sure of the best way to design this.  It might be best to have a merge function like above or perhaps just a synchronize names function that is orthogonal to merging.



Discussion:
On second thought, I probably do need the two dasks separately, so synchronizing names without a simultaneous merge is likely of value.

This should also skip gracefully in cases where the values can not be compared.


Resolving patch:
diff --git a/dask/compatibility.py b/dask/compatibility.py
--- a/dask/compatibility.py
+++ b/dask/compatibility.py
@@ -8,6 +8,7 @@
 if PY3:
     import builtins
     from queue import Queue
+    from itertools import zip_longest
     unicode = str
     long = int
     def apply(func, args):
@@ -16,6 +17,7 @@ def apply(func, args):
     import __builtin__ as builtins
     from Queue import Queue
     import operator
+    from itertools import izip_longest as zip_longest
     unicode = unicode
     long = long
     apply = apply
diff --git a/dask/core.py b/dask/core.py
--- a/dask/core.py
+++ b/dask/core.py
@@ -41,6 +41,21 @@ def istask(x):
     return isinstance(x, tuple) and x and callable(x[0])
 
 
+def preorder_traversal(task):
+    """A generator to preorder-traverse a task."""
+
+    for item in task:
+        if istask(item):
+            for i in preorder_traversal(item):
+                yield i
+        elif isinstance(item, list):
+            yield list
+            for i in preorder_traversal(item):
+                yield i
+        else:
+            yield item
+
+
 def _get_task(d, task, maxdepth=1000):
     # non-recursive.  DAG property is checked upon reaching maxdepth.
     _iter = lambda *args: iter(args)
@@ -382,4 +397,3 @@ def isdag(d, keys):
     getcycle
     """
     return not getcycle(d, keys)
-
diff --git a/dask/optimize.py b/dask/optimize.py
--- a/dask/optimize.py
+++ b/dask/optimize.py
@@ -1,5 +1,8 @@
+from itertools import count
+from .compatibility import zip_longest
 from .core import (istask, get_dependencies, subs, toposort, flatten,
-                   reverse_dict, add, inc, ishashable)
+                   reverse_dict, add, inc, ishashable, preorder_traversal)
+from .rewrite import END
 from toolz import identity
 
 
@@ -242,3 +245,198 @@ def dealias(dsk):
         else:
             dsk3[k] = (identity, k2)
     return dsk3
+
+
+def equivalent(term1, term2, subs=None):
+    """Determine if two terms are equivalent, modulo variable substitution.
+
+    Equivalent to applying substitutions in `subs` to `term2`, then checking if
+    `term1 == term2`.
+
+    If a subterm doesn't support comparison (i.e. `term1 == term2` errors),
+    returns `False`.
+
+    Parameters
+    ----------
+    term1, term2 : terms
+    subs : dict, optional
+        Mapping of substitutions from `term2` to `term1`
+
+    Example
+    -------
+    >>> from operator import add
+    >>> term1 = (add, 'a', 'b')
+    >>> term2 = (add, 'x', 'y')
+    >>> subs = {'x': 'a', 'y': 'b'}
+    >>> equivalent(term1, term2, subs)
+    True
+    >>> subs = {'x': 'a'}
+    >>> equivalent(term1, term2, subs)
+    False
+    """
+
+    # Quick escape for special cases
+    head_type = type(term1)
+    if type(term2) != head_type:
+        # If terms aren't same type, fail
+        return False
+    elif head_type not in (tuple, list):
+        # For literals, just compare
+        try:
+            # `is` is tried first, to allow objects that don't implement `==`
+            # to work for cases where term1 is term2. If `is` returns False,
+            # and `==` errors, then the only thing we can do is return False.
+            return term1 is term2 or term1 == term2
+        except:
+            return False
+
+    pot1 = preorder_traversal(term1)
+    pot2 = preorder_traversal(term2)
+    subs = {} if subs is None else subs
+
+    for t1, t2 in zip_longest(pot1, pot2, fillvalue=END):
+        if t1 is END or t2 is END:
+            # If terms aren't same length: fail
+            return False
+        elif ishashable(t2) and t2 in subs:
+            val = subs[t2]
+        else:
+            val = t2
+        try:
+            if t1 is not t2 and t1 != val:
+                return False
+        except:
+            return False
+    return True
+
+
+def dependency_dict(dsk):
+    """Create a dict matching ordered dependencies to keys.
+
+    Example
+    -------
+    >>> from operator import add
+    >>> dsk = {'a': 1, 'b': 2, 'c': (add, 'a', 'a'), 'd': (add, 'b', 'a')}
+    >>> dependency_dict(dsk)    # doctest: +SKIP
+    {(): ['a', 'b'], ('a', 'a'): ['c'], ('b', 'a'): ['d']}
+    """
+
+    dep_dict = {}
+    for key in dsk:
+        deps = tuple(get_dependencies(dsk, key, True))
+        dep_dict.setdefault(deps, []).append(key)
+    return dep_dict
+
+
+def _possible_matches(dep_dict, deps, subs):
+    deps2 = []
+    for d in deps:
+        v = subs.get(d, None)
+        if v is not None:
+            deps2.append(v)
+        else:
+            return []
+    deps2 = tuple(deps2)
+    return dep_dict.get(deps2, [])
+
+
+def _sync_keys(dsk1, dsk2, dsk2_topo):
+    dep_dict1 = dependency_dict(dsk1)
+    subs = {}
+
+    for key2 in toposort(dsk2):
+        deps = tuple(get_dependencies(dsk2, key2, True))
+        # List of keys in dsk1 that have terms that *may* match key2
+        possible_matches = _possible_matches(dep_dict1, deps, subs)
+        if possible_matches:
+            val2 = dsk2[key2]
+            for key1 in possible_matches:
+                val1 = dsk1[key1]
+                if equivalent(val1, val2, subs):
+                    subs[key2] = key1
+                    break
+    return subs
+
+
+def sync_keys(dsk1, dsk2):
+    """Return a dict matching keys in `dsk2` to equivalent keys in `dsk1`.
+
+    Parameters
+    ----------
+    dsk1, dsk2 : dict
+
+    Example
+    -------
+    >>> from operator import add, mul
+    >>> dsk1 = {'a': 1, 'b': (add, 'a', 10), 'c': (mul, 'b', 5)}
+    >>> dsk2 = {'x': 1, 'y': (add, 'x', 10), 'z': (mul, 'y', 2)}
+    >>> sync_keys(dsk1, dsk2)   # doctest: +SKIP
+    {'x': 'a', 'y': 'b'}
+    """
+
+    return _sync_keys(dsk1, dsk2, toposort(dsk2))
+
+
+def merge_sync(dsk1, dsk2):
+    """Merge two dasks together, combining equivalent tasks.
+
+    If a task in `dsk2` exists in `dsk1`, the task and key from `dsk1` is used.
+    If a task in `dsk2` has the same key as a task in `dsk1` (and they aren't
+    equivalent tasks), then a new key is created for the task in `dsk2`. This
+    prevents name conflicts.
+
+    Parameters
+    ----------
+    dsk1, dsk2 : dict
+        Variable names in `dsk2` are replaced with equivalent ones in `dsk1`
+        before merging.
+
+    Returns
+    -------
+    new_dsk : dict
+        The merged dask.
+    key_map : dict
+        A mapping between the keys from `dsk2` to their new names in `new_dsk`.
+
+    Example
+    -------
+    >>> from operator import add, mul
+    >>> dsk1 = {'a': 1, 'b': (add, 'a', 10), 'c': (mul, 'b', 5)}
+    >>> dsk2 = {'x': 1, 'y': (add, 'x', 10), 'z': (mul, 'y', 2)}
+    >>> new_dsk, key_map = merge_sync(dsk1, dsk2)
+    >>> new_dsk     # doctest: +SKIP
+    {'a': 1, 'b': (add, 'a', 10), 'c': (mul, 'b', 5), 'z': (mul, 'b', 2)}
+    >>> key_map     # doctest: +SKIP
+    {'x': 'a', 'y': 'b', 'z': 'z'}
+
+    Conflicting names are replaced with auto-generated names upon merging.
+
+    >>> dsk1 = {'a': 1, 'res': (add, 'a', 1)}
+    >>> dsk2 = {'x': 1, 'res': (add, 'x', 2)}
+    >>> new_dsk, key_map = merge_sync(dsk1, dsk2)
+    >>> new_dsk     # doctest: +SKIP
+    {'a': 1, 'res': (add, 'a', 1), 'merge_1': (add, 'a', 2)}
+    >>> key_map     # doctest: +SKIP
+    {'x': 'a', 'res': 'merge_1'}
+    """
+
+    dsk2_topo = toposort(dsk2)
+    sd = _sync_keys(dsk1, dsk2, dsk2_topo)
+    new_dsk = dsk1.copy()
+    for key in dsk2_topo:
+        if key in sd:
+            new_key = sd[key]
+        else:
+            if key in dsk1:
+                new_key = next(merge_sync.names)
+            else:
+                new_key = key
+            sd[key] = new_key
+        task = dsk2[key]
+        for a, b in sd.items():
+            task = subs(task, a, b)
+        new_dsk[new_key] = task
+    return new_dsk, sd
+
+# store the name iterator in the function
+merge_sync.names = ('merge_%d' % i for i in count(1))


Test patch:
diff --git a/dask/tests/test_core.py b/dask/tests/test_core.py
--- a/dask/tests/test_core.py
+++ b/dask/tests/test_core.py
@@ -1,5 +1,6 @@
 from dask.utils import raises
-from dask.core import istask, get, get_dependencies, flatten, subs
+from dask.core import (istask, get, get_dependencies, flatten, subs,
+        preorder_traversal)
 
 
 def contains(a, b):
@@ -30,6 +31,15 @@ def test_istask():
      ':z': (add, ':x', ':y')}
 
 
+def test_preorder_traversal():
+    t = (add, 1, 2)
+    assert list(preorder_traversal(t)) == [add, 1, 2]
+    t = (add, (add, 1, 2), (add, 3, 4))
+    assert list(preorder_traversal(t)) == [add, add, 1, 2, add, 3, 4]
+    t = (add, (sum, [1, 2]), 3)
+    assert list(preorder_traversal(t)) == [add, sum, list, 1, 2, 3]
+
+
 def test_get():
     assert get(d, ':x') == 1
     assert get(d, ':y') == 2
diff --git a/dask/tests/test_optimize.py b/dask/tests/test_optimize.py
--- a/dask/tests/test_optimize.py
+++ b/dask/tests/test_optimize.py
@@ -1,14 +1,14 @@
+from itertools import count
+from operator import add, mul
 from toolz import partial, identity
 from dask.utils import raises
 from dask.optimize import (cull, fuse, inline, inline_functions, functions_of,
-        dealias)
+        dealias, equivalent, sync_keys, merge_sync)
 
 
 def inc(x):
     return x + 1
 
-def add(x, y):
-    return x + y
 
 def double(x):
     return x * 2
@@ -186,3 +186,109 @@ def test_dealias():
                 'c': (identity, 'a')}
 
     assert dealias(dsk)  == expected
+
+
+def test_equivalent():
+    t1 = (add, 'a', 'b')
+    t2 = (add, 'x', 'y')
+
+    assert equivalent(t1, t1)
+    assert not equivalent(t1, t2)
+    assert equivalent(t1, t2, {'x': 'a', 'y': 'b'})
+    assert not equivalent(t1, t2, {'a': 'x'})
+
+    t1 = (add, (double, 'a'), (double, 'a'))
+    t2 = (add, (double, 'b'), (double, 'c'))
+
+    assert equivalent(t1, t1)
+    assert not equivalent(t1, t2)
+    assert equivalent(t1, t2, {'b': 'a', 'c': 'a'})
+    assert not equivalent(t1, t2, {'b': 'a', 'c': 'd'})
+    assert not equivalent(t2, t1, {'a': 'b'})
+
+    # Test literal comparisons
+    assert equivalent(1, 1)
+    assert not equivalent(1, 2)
+    assert equivalent((1, 2, 3), (1, 2, 3))
+
+
+class Uncomparable(object):
+    def __eq__(self, other):
+        raise TypeError("Uncomparable type")
+
+
+def test_equivalence_uncomparable():
+    t1 = Uncomparable()
+    t2 = Uncomparable()
+    assert raises(TypeError, lambda: t1 == t2)
+    assert equivalent(t1, t1)
+    assert not equivalent(t1, t2)
+    assert equivalent((add, t1, 0), (add, t1, 0))
+    assert not equivalent((add, t1, 0), (add, t2, 0))
+
+
+def test_sync_keys():
+    dsk1 = {'a': 1, 'b': (add, 'a', 10), 'c': (mul, 'b', 5)}
+    dsk2 = {'x': 1, 'y': (add, 'x', 10), 'z': (mul, 'y', 2)}
+    assert sync_keys(dsk1, dsk2) == {'x': 'a', 'y': 'b'}
+    assert sync_keys(dsk2, dsk1) == {'a': 'x', 'b': 'y'}
+
+    dsk1 = {'a': 1, 'b': 2, 'c': (add, 'a', 'b'), 'd': (inc, (add, 'a', 'b'))}
+    dsk2 = {'x': 1, 'y': 5, 'z': (add, 'x', 'y'), 'w': (inc, (add, 'x', 'y'))}
+    assert sync_keys(dsk1, dsk2) == {'x': 'a'}
+    assert sync_keys(dsk2, dsk1) == {'a': 'x'}
+
+
+def test_sync_uncomparable():
+    t1 = Uncomparable()
+    t2 = Uncomparable()
+    dsk1 = {'a': 1, 'b': t1, 'c': (add, 'a', 'b')}
+    dsk2 = {'x': 1, 'y': t2, 'z': (add, 'y', 'x')}
+    assert sync_keys(dsk1, dsk2) == {'x': 'a'}
+
+    dsk2 = {'x': 1, 'y': t1, 'z': (add, 'y', 'x')}
+    assert sync_keys(dsk1, dsk2) == {'x': 'a', 'y': 'b'}
+
+
+def test_merge_sync():
+    dsk1 = {'a': 1, 'b': (add, 'a', 10), 'c': (mul, 'b', 5)}
+    dsk2 = {'x': 1, 'y': (add, 'x', 10), 'z': (mul, 'y', 2)}
+    new_dsk, key_map = merge_sync(dsk1, dsk2)
+    assert new_dsk == {'a': 1, 'b': (add, 'a', 10), 'c': (mul, 'b', 5),
+            'z': (mul, 'b', 2)}
+    assert key_map == {'x': 'a', 'y': 'b', 'z': 'z'}
+
+    dsk1 = {'g1': 1,
+            'g2': 2,
+            'g3': (add, 'g1', 1),
+            'g4': (add, 'g2', 1),
+            'g5': (mul, (inc, 'g3'), (inc, 'g4'))}
+    dsk2 = {'h1': 1,
+            'h2': 5,
+            'h3': (add, 'h1', 1),
+            'h4': (add, 'h2', 1),
+            'h5': (mul, (inc, 'h3'), (inc, 'h4'))}
+    new_dsk, key_map = merge_sync(dsk1, dsk2)
+    assert new_dsk == {'g1': 1,
+                       'g2': 2,
+                       'g3': (add, 'g1', 1),
+                       'g4': (add, 'g2', 1),
+                       'g5': (mul, (inc, 'g3'), (inc, 'g4')),
+                       'h2': 5,
+                       'h4': (add, 'h2', 1),
+                       'h5': (mul, (inc, 'g3'), (inc, 'h4'))}
+    assert key_map == {'h1': 'g1', 'h2': 'h2', 'h3': 'g3',
+            'h4': 'h4', 'h5': 'h5'}
+
+    # Test merging with name conflict
+    # Reset name count to ensure same numbers
+    merge_sync.names = ("merge_%d" % i for i in count(1))
+    dsk1 = {'g1': 1, 'conflict': (add, 'g1', 2), 'g2': (add, 'conflict', 3)}
+    dsk2 = {'h1': 1, 'conflict': (add, 'h1', 4), 'h2': (add, 'conflict', 3)}
+    new_dsk, key_map = merge_sync(dsk1, dsk2)
+    assert new_dsk == {'g1': 1,
+                       'conflict': (add, 'g1', 2),
+                       'merge_1': (add, 'g1', 4),
+                       'g2': (add, 'conflict', 3),
+                       'h2': (add, 'merge_1', 3)}
+    assert key_map == {'h1': 'g1', 'conflict': 'merge_1', 'h2': 'h2'}