Skip empty partitions when doing groupby.value_counts (dask#7073)

jsignell · web-flow · commit 30c9a3a34240 · 2021-01-20T16:35:54.000-06:00
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
@@ -1931,7 +1931,7 @@ def agg(self, arg, split_every=None, split_out=1):
     def value_counts(self, split_every=None, split_out=1):
         return self._aca_agg(
             token="value_counts",
-            func=M.value_counts,
+            func=_value_counts,
             aggfunc=_value_counts_aggregate,
             split_every=split_every,
             split_out=split_out,
@@ -1956,6 +1956,13 @@ def _unique_aggregate(series_gb, name=None):
     return ret
 
 
+def _value_counts(x, **kwargs):
+    if len(x):
+        return M.value_counts(x, **kwargs)
+    else:
+        return pd.Series(dtype=int)
+
+
 def _value_counts_aggregate(series_gb):
     to_concat = {k: v.sum(level=1) for k, v in series_gb}
     names = list(series_gb.obj.index.names)
diff --git a/dask/dataframe/tests/test_groupby.py b/dask/dataframe/tests/test_groupby.py
@@ -2475,3 +2475,26 @@ def agg(grp, **kwargs):
         agg(pdf.groupby(groupby, observed=observed)),
         agg(ddf.groupby(groupby, observed=observed)),
     )
+
+
+def test_empty_partitions_with_value_counts():
+    # https://github.com/dask/dask/issues/7065
+    df = pd.DataFrame(
+        data=[
+            ["a1", "b1"],
+            ["a1", None],
+            ["a1", "b1"],
+            [None, None],
+            [None, None],
+            [None, None],
+            ["a3", "b3"],
+            ["a3", "b3"],
+            ["a5", "b5"],
+        ],
+        columns=["A", "B"],
+    )
+
+    expected = df.groupby("A")["B"].value_counts()
+    ddf = dd.from_pandas(df, npartitions=3)
+    actual = ddf.groupby("A")["B"].value_counts()
+    assert_eq(expected, actual)
