Repository: dask/dask

Issue:
SeriesGroupBy.nunique may output incorrect result
```
dsk = {('x', 0): pd.DataFrame({'a': [1, 2, 6], 
                               'b': [4, 2., 7]}, index=[0, 1, 3]),
       ('x', 1): pd.DataFrame({'a': [4, 2, 6],
                               'b': [3, 3, 1]}, index=[5, 6, 8]),
       ('x', 2): pd.DataFrame({'a': [4, 3, 7],
                               'b': [1, 1, 3]}, index=[9, 9, 9])}
d = dd.DataFrame(dsk, 'x', ['a', 'b'], [0, 4, 9, 9])
full = d.compute()

d.groupby(d.b > 3).a.nunique().compute()
# b
# False    6
# True     2
# Name: a, dtype: int64

full.groupby(full.b > 3).a.nunique()
# b
# False    5
# True     2
# Name: a, dtype: int64
```



Discussion:

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -1088,9 +1088,10 @@ def nunique(self):
         def chunk(df, index):
             # we call set_index here to force a possibly duplicate index
             # for our reduce step
-            return (df.groupby(index)
-                      .apply(pd.DataFrame.drop_duplicates, subset=self.key)
-                      .set_index(index))
+            grouped = (df.groupby(index)
+                .apply(pd.DataFrame.drop_duplicates, subset=self.key))
+            grouped.index = grouped.index.get_level_values(level=0)
+            return grouped
 
         def agg(df):
             return df.groupby(level=0)[self.key].nunique()


Test patch:
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -162,13 +162,18 @@ def test_split_apply_combine_on_series():
         assert eq(d.groupby(ddkey).a.max(), full.groupby(pdkey).a.max())
         assert eq(d.groupby(ddkey).a.count(), full.groupby(pdkey).a.count())
 
-    for ddkey, pdkey in [(d.b > 3, full.b > 3)]:
-        # temp
-        assert eq(d.groupby(ddkey).a.sum(), full.groupby(pdkey).a.sum())
-        assert eq(d.groupby(ddkey).a.min(), full.groupby(pdkey).a.min())
-        assert eq(d.groupby(ddkey).a.max(), full.groupby(pdkey).a.max())
-        assert eq(d.groupby(ddkey).a.count(), full.groupby(pdkey).a.count())
-        # assert eq(d.groupby(ddkey).a.nunique(), full.groupby(pdkey).a.nunique())
+    for i in range(8):
+        assert eq(d.groupby(d.b > i).a.sum(), full.groupby(full.b > i).a.sum())
+        assert eq(d.groupby(d.b > i).a.min(), full.groupby(full.b > i).a.min())
+        assert eq(d.groupby(d.b > i).a.max(), full.groupby(full.b > i).a.max())
+        assert eq(d.groupby(d.b > i).a.count(), full.groupby(full.b > i).a.count())
+        assert eq(d.groupby(d.b > i).a.nunique(), full.groupby(full.b > i).a.nunique())
+
+        assert eq(d.groupby(d.a > i).b.sum(), full.groupby(full.a > i).b.sum())
+        assert eq(d.groupby(d.a > i).b.min(), full.groupby(full.a > i).b.min())
+        assert eq(d.groupby(d.a > i).b.max(), full.groupby(full.a > i).b.max())
+        assert eq(d.groupby(d.a > i).b.count(), full.groupby(full.a > i).b.count())
+        assert eq(d.groupby(d.a > i).b.nunique(), full.groupby(full.a > i).b.nunique())
 
     for ddkey, pdkey in [('a', 'a'), (d.a, full.a),
                           (d.a + 1, full.a + 1), (d.a > 3, full.a > 3)]: