Repository: dask/dask

Issue:
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


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`.

Resolving patch:
diff --git a/dask/bag/core.py b/dask/bag/core.py
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -19,13 +19,13 @@
 try:
     import cytoolz
     from cytoolz import (frequencies, merge_with, join, reduceby,
-                         count, pluck, groupby, topk)
+                         count, pluck, groupby, topk, unique)
     if LooseVersion(cytoolz.__version__) > '0.7.3':
         from cytoolz import accumulate  # noqa: F811
         _implement_accumulate = True
 except ImportError:
     from toolz import (frequencies, merge_with, join, reduceby,
-                       count, pluck, groupby, topk)
+                       count, pluck, groupby, topk, unique)
 
 from .. import config
 from .avro import to_avro
@@ -817,17 +817,31 @@ 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,
-                              name='distinct')
+        func = chunk_distinct if key is None else partial(chunk_distinct, key=key)
+        agg = merge_distinct if key is None else partial(merge_distinct, key=key)
+        return self.reduction(func, agg, out_type=Bag, name='distinct')
 
     def reduction(self, perpartition, aggregate, split_every=None,
                   out_type=Item, name=None):
@@ -1630,8 +1644,15 @@ def from_delayed(values):
     return Bag(graph, name, len(values))
 
 
-def merge_distinct(seqs):
-    return set().union(*seqs)
+def chunk_distinct(seq, key=None):
+    key2 = key
+    if key is not None and not callable(key):
+        key2 = lambda x: x[key]
+    return list(unique(seq, key=key2))
+
+
+def merge_distinct(seqs, key=None):
+    return chunk_distinct(toolz.concat(seqs), key=key)
 
 
 def merge_frequencies(seqs):


Test patch:
diff --git a/dask/bag/tests/test_bag.py b/dask/bag/tests/test_bag.py
--- a/dask/bag/tests/test_bag.py
+++ b/dask/bag/tests/test_bag.py
@@ -8,7 +8,7 @@
 
 import partd
 import pytest
-from toolz import merge, join, filter, identity, valmap, groupby, pluck
+from toolz import merge, join, filter, identity, valmap, groupby, pluck, unique
 
 import dask
 import dask.bag as db
@@ -227,6 +227,14 @@ def test_distinct():
     assert bag.filter(None).distinct().compute() == []
 
 
+def test_distinct_with_key():
+    seq = [{'a': i} for i in [0, 1, 2, 1, 2, 3, 2, 3, 4, 5]]
+    bag = db.from_sequence(seq, npartitions=3)
+    expected = list(unique(seq, key=lambda x: x['a']))
+    assert_eq(bag.distinct(key='a'), expected)
+    assert_eq(bag.distinct(key=lambda x: x['a']), expected)
+
+
 def test_frequencies():
     c = b.frequencies()
     assert dict(c) == {0: 3, 1: 3, 2: 3, 3: 3, 4: 3}