Repository: dask/dask
Pull request: Add key argument to Bag.distinct
- [x] Tests added / passed
- [x] Passes `flake8 dask`

closes #2493 

A few things to discuss:

1. Should we change the implementation to use `fold` instead of `reduction`?
2. Is it ok to default `key` to `toolz.identity` instead of `None` + conditional logic?

Problem statement:
Ability to drop duplicates on Bag based on subset of data
It would be nice if it was possible to drop duplicates on a bag based on a subset of the data, similar to `DataFrame.drop_duplicates(subset=[...])`.  Perhaps the API could be `Bag.distinct(key=lambda x: x['name'])`.

I can help implement this feature, just let me know what the ideal API would be & any ideas on making it performant are greatly appreciated. :D

Hints from issue discussion:
Hi @munro this seems useful to me.  The current thing to look at is (as you seem to have found) [Bag.distinct](http://dask.pydata.org/en/latest/bag-api.html#dask.bag.Bag.distinct).  This operates as a reduction, and assumes that the output set will fit nicely into memory in a single partition.  Adding a `key=` keyword seems sensible if it is likely to be used in practice.  Have you taken a look at the current implementation?  Do you have any particular questions?

In case you haven't seen it [this page](http://dask.pydata.org/en/latest/develop.html) has notes about style and testing and such.
I think following the behavior of `toolz.unique` would be good.

It also might be nice to use the same function name here (see https://github.com/pytoolz/toolz/issues/276). Could either deprecate `Bag.distinct` to `Bag.unique`, or `toolz.unique` to `toolz.distinct`.
I'll give this a try!

I gave the implementation a gander and it looks like I should edit [this](https://github.com/dask/dask/blob/master/dask/bag/core.py#L820-L830) to support a `key` argument and use `toolz.unique` for implementation with tests in `dask/bag/tests/test_bag.py`.

Commit message: Removed unused import

Patch to review:
diff --git a/dask/bag/core.py b/dask/bag/core.py
index 99f5645ecdb..3a9e923251a 100644
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -817,16 +817,34 @@ def topk(self, k, key=None, split_every=None):
         return self.reduction(func, compose(func, toolz.concat), out_type=Bag,
                               split_every=split_every, name='topk')
 
-    def distinct(self):
+    def distinct(self, key=None):
         """ Distinct elements of collection
 
         Unordered without repeats.
 
+        Parameters
+        ----------
+        key: {callable,str}
+            Defines uniqueness of items in bag by calling `key` on each item.
+            If a string is passed `key` is considered to be `lambda x: x[key]`.
+
+        Examples
+        --------
         >>> b = from_sequence(['Alice', 'Bob', 'Alice'])
         >>> sorted(b.distinct())
         ['Alice', 'Bob']
+        >>> b = from_sequence([{'name': 'Alice'}, {'name': 'Bob'}, {'name': 'Alice'}])
+        >>> b.distinct(key=lambda x: x['name']).compute()
+        [{'name': 'Alice'}, {'name': 'Bob'}]
+        >>> b.distinct(key='name').compute()
+        [{'name': 'Alice'}, {'name': 'Bob'}]
         """
-        return self.reduction(set, merge_distinct, out_type=Bag,
+        key_func = key if callable(key) or key is None else lambda x: x[key]
+        perpartition = lambda seq: list(toolz.unique(seq, key=key_func))
+        aggregate = merge_distinct(key=key_func)
+        return self.reduction(perpartition,
+                              aggregate,
+                              out_type=Bag,
                               name='distinct')
 
     def reduction(self, perpartition, aggregate, split_every=None,
@@ -1630,8 +1648,14 @@ def from_delayed(values):
     return Bag(graph, name, len(values))
 
 
-def merge_distinct(seqs):
-    return set().union(*seqs)
+@curry
+def merge_distinct(seqs, key=None):
+    if key is None:
+        return list(toolz.unique(toolz.concat(seqs)))
+    elif callable(key):
+        return list(toolz.unique(toolz.concat(seqs), key=key))
+    else:
+        return list(toolz.unique(toolz.concat(seqs), key=lambda x: x[key]))
 
 
 def merge_frequencies(seqs):
diff --git a/dask/bag/tests/test_bag.py b/dask/bag/tests/test_bag.py
index 01a0e3febc0..251e3e47076 100644
--- a/dask/bag/tests/test_bag.py
+++ b/dask/bag/tests/test_bag.py
@@ -14,7 +14,7 @@
 import dask.bag as db
 from dask.bag.core import (Bag, lazify, lazify_task, map, collect,
                            reduceby, reify, partition, inline_singleton_lists,
-                           optimize, from_delayed)
+                           optimize, from_delayed, merge_distinct)
 from dask.bag.utils import assert_eq
 from dask.compatibility import BZ2File, GzipFile, PY2, Iterator
 from dask.delayed import Delayed
@@ -227,6 +227,24 @@ def test_distinct():
     assert bag.filter(None).distinct().compute() == []
 
 
+def test_distinct_with_key():
+    bag = db.from_sequence([{'a': 0}, {'a': 0}, {'a': 1}])
+    expected = [{'a': 0}, {'a': 1}]
+    key = 'a'
+    key_func = lambda x: x[key]
+    assert_eq(bag.distinct(key_func), expected)
+    assert_eq(bag.distinct(key), expected)
+
+
+def test_merge_distinct():
+    a = [[0, 1], [1, 2]]
+    b = [[{'a': 0}, {'a': 1}], [{'a': 1}, {'a': 2}]]
+    key = 'a'
+    key_func = lambda x: x[key]
+    assert merge_distinct(a) == [0, 1, 2]
+    assert merge_distinct(b, key_func) == [{'a': 0}, {'a': 1}, {'a': 2}]
+
+
 def test_frequencies():
     c = b.frequencies()
     assert dict(c) == {0: 3, 1: 3, 2: 3, 3: 3, 4: 3}


Code review comments:
diff_hunk: @@ -817,16 +817,34 @@ def topk(self, k, key=None, split_every=None):
         return self.reduction(func, compose(func, toolz.concat), out_type=Bag,
                               split_every=split_every, name='topk')
 
-    def distinct(self):
+    def distinct(self, key=None):
         """ Distinct elements of collection
 
         Unordered without repeats.
 
+        Parameters
+        ----------
+        key: {callable,str}
+            Defines uniqueness of items in bag by calling `key` on each item.
+            If a string is passed `key` is considered to be `lambda x: x[key]`.
+
+        Examples
+        --------
         >>> b = from_sequence(['Alice', 'Bob', 'Alice'])
         >>> sorted(b.distinct())
         ['Alice', 'Bob']
+        >>> b = from_sequence([{'name': 'Alice'}, {'name': 'Bob'}, {'name': 'Alice'}])
+        >>> b.distinct(key=lambda x: x['name']).compute()
+        [{'name': 'Alice'}, {'name': 'Bob'}]
+        >>> b.distinct(key='name').compute()
+        [{'name': 'Alice'}, {'name': 'Bob'}]
         """
-        return self.reduction(set, merge_distinct, out_type=Bag,
+        key_func = key if callable(key) or key is None else lambda x: x[key]
path: dask/bag/core.py
text: @user1:
Similarly I would handle this within the `perpartition` or `merge_distinct` functions in order to avoid passing around lambdas in the graph.

diff_hunk: @@ -817,16 +817,34 @@ def topk(self, k, key=None, split_every=None):
         return self.reduction(func, compose(func, toolz.concat), out_type=Bag,
                               split_every=split_every, name='topk')
 
-    def distinct(self):
+    def distinct(self, key=None):
         """ Distinct elements of collection
 
         Unordered without repeats.
 
+        Parameters
+        ----------
+        key: {callable,str}
+            Defines uniqueness of items in bag by calling `key` on each item.
+            If a string is passed `key` is considered to be `lambda x: x[key]`.
+
+        Examples
+        --------
         >>> b = from_sequence(['Alice', 'Bob', 'Alice'])
         >>> sorted(b.distinct())
         ['Alice', 'Bob']
+        >>> b = from_sequence([{'name': 'Alice'}, {'name': 'Bob'}, {'name': 'Alice'}])
+        >>> b.distinct(key=lambda x: x['name']).compute()
+        [{'name': 'Alice'}, {'name': 'Bob'}]
+        >>> b.distinct(key='name').compute()
+        [{'name': 'Alice'}, {'name': 'Bob'}]
         """
-        return self.reduction(set, merge_distinct, out_type=Bag,
+        key_func = key if callable(key) or key is None else lambda x: x[key]
+        perpartition = lambda seq: list(toolz.unique(seq, key=key_func))
path: dask/bag/core.py
text: @user1:
We'd like to avoid the dynamic creation of functions here.  They are somewhat costly to serialize.  I recommend making another top level function like `merge_distinct`

diff_hunk: @@ -1630,8 +1648,14 @@ def from_delayed(values):
     return Bag(graph, name, len(values))
 
 
-def merge_distinct(seqs):
-    return set().union(*seqs)
+@curry
path: dask/bag/core.py
text: @user1:
It would be good to avoid `curry` here if possible.  Currying can easily generate complex situations.  I recommend either using `functools.partial` or passing keyword arguments explicitly (this would require changes to `reduction` to optionally support new `perpartition_kwargs` and `aggregate_kwargs` keyword arguments.  Short term I recommend going with partial because it's probably easier.

diff_hunk: @@ -1630,8 +1648,14 @@ def from_delayed(values):
     return Bag(graph, name, len(values))
 
 
-def merge_distinct(seqs):
-    return set().union(*seqs)
+@curry
+def merge_distinct(seqs, key=None):
+    if key is None:
+        return list(toolz.unique(toolz.concat(seqs)))
path: dask/bag/core.py
text: @user1:
If cytoolz is around then we'll probably want to use cytoolz.unique instead.  See the toolz and cytoolz imports at the top of this file.