Add support for normalize in value_counts (#7342)

dask/dataframe/core.py

@@ -3122,13 +3122,20 @@ def nunique(self, split_every=None):
 
     @derived_from(pd.Series)
     def value_counts(
-        self, sort=None, ascending=False, dropna=None, split_every=None, split_out=1
+        self,
+        sort=None,
+        ascending=False,
+        dropna=None,
+        normalize=False,
+        split_every=None,
+        split_out=1,
     ):
         """
         Note: dropna is only supported in pandas >= 1.1.0, in which case it defaults to
         True.
         """
         kwargs = {"sort": sort, "ascending": ascending}
+
         if dropna is not None:
             if not PANDAS_GT_110:
                 raise NotImplementedError(
@@ -3137,16 +3144,23 @@ def value_counts(
                 )
             kwargs["dropna"] = dropna
 
+        aggregate_kwargs = {"normalize": normalize}
+        if split_out > 1:
+            aggregate_kwargs["total_length"] = (
+                len(self) if dropna is False else len(self.dropna())
+            )
+
         return aca(
             self,
             chunk=M.value_counts,
             aggregate=methods.value_counts_aggregate,
             combine=methods.value_counts_combine,
-            meta=self._meta.value_counts(),
+            meta=self._meta.value_counts(normalize=normalize),
             token="value-counts",
             split_every=split_every,
             split_out=split_out,
             split_out_setup=split_out_on_index,
+            aggregate_kwargs=aggregate_kwargs,
             **kwargs,
         )
 

dask/dataframe/methods.py

@@ -309,8 +309,12 @@ def value_counts_combine(x, sort=True, ascending=False, **groupby_kwargs):
     return x.groupby(level=0, **groupby_kwargs).sum()
 
 
-def value_counts_aggregate(x, sort=True, ascending=False, **groupby_kwargs):
+def value_counts_aggregate(
+    x, sort=True, ascending=False, normalize=False, total_length=None, **groupby_kwargs
+):
     out = value_counts_combine(x, **groupby_kwargs)
+    if normalize:
+        out /= total_length if total_length is not None else out.sum()
     if sort:
         return out.sort_values(ascending=ascending)
     return out

dask/dataframe/tests/test_dataframe.py

@@ -1079,6 +1079,44 @@ def test_value_counts_with_dropna():
     assert result._name != result2._name
 
 
+def test_value_counts_with_normalize():
+    df = pd.DataFrame({"x": [1, 2, 1, 3, 3, 1, 4]})
+    ddf = dd.from_pandas(df, npartitions=3)
+    result = ddf.x.value_counts(normalize=True)
+    expected = df.x.value_counts(normalize=True)
+    assert_eq(result, expected)
+
+    result2 = ddf.x.value_counts(split_every=2, normalize=True)
+    assert_eq(result2, expected)
+    assert result._name != result2._name
+
+    result3 = ddf.x.value_counts(split_out=2, normalize=True)
+    assert_eq(result3, expected)
+    assert result._name != result3._name
+
+
+@pytest.mark.skipif(not PANDAS_GT_110, reason="dropna implemented in pandas 1.1.0")
+def test_value_counts_with_normalize_and_dropna():
+    df = pd.DataFrame({"x": [1, 2, 1, 3, np.nan, 1, 4]})
+    ddf = dd.from_pandas(df, npartitions=3)
+
+    result = ddf.x.value_counts(dropna=False, normalize=True)
+    expected = df.x.value_counts(dropna=False, normalize=True)
+    assert_eq(result, expected)
+
+    result2 = ddf.x.value_counts(split_every=2, dropna=False, normalize=True)
+    assert_eq(result2, expected)
+    assert result._name != result2._name
+
+    result3 = ddf.x.value_counts(split_out=2, dropna=False, normalize=True)
+    assert_eq(result3, expected)
+    assert result._name != result3._name
+
+    result4 = ddf.x.value_counts(dropna=True, normalize=True, split_out=2)
+    expected4 = df.x.value_counts(dropna=True, normalize=True)
+    assert_eq(result4, expected4)
+
+
 def test_unique():
     pdf = pd.DataFrame(
         {