Repository: dask/dask

Issue:
Add DataFrame.nunique()
Opening this as a new issue, as it came up in a recent [SO post](https://stackoverflow.com/questions/69928786/number-of-unique-values-in-dask-dataframe-columns) that @ian-r-rose and I worked on. It seems like it was mistakenly marked as added in issue #3817.


Discussion:
@scharlottej13 can you please add a bit more information to the description? A sentence or two summarising the SO post and/or issue if that it relevant.
I think the SO post also mentions `groupby.nunique`. is this missing as well and part of this issue?
This is to write the Dask version of [pandas.DataFrame.nunique](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.nunique.html). The linked SO post is one example of someone asking for this functionality. I don't think `groupby.nunique` is missing, I found the function [here](https://github.com/dask/dask/blob/main/dask/dataframe/groupby.py#L1940), so that's not part of this issue.

@fjetter Let me know if we need any other details!

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -3489,8 +3489,13 @@ def unique(self, split_every=None, split_out=1):
         )
 
     @derived_from(pd.Series)
-    def nunique(self, split_every=None):
-        return self.drop_duplicates(split_every=split_every).count()
+    def nunique(self, split_every=None, dropna=True):
+        uniqs = self.drop_duplicates(split_every=split_every)
+        if dropna:
+            # count mimics pandas behavior and excludes NA values
+            return uniqs.count()
+        else:
+            return uniqs.size
 
     @derived_from(pd.Series)
     def value_counts(
@@ -5129,6 +5134,38 @@ def applymap(self, func, meta="__no_default__"):
     def round(self, decimals=0):
         return elemwise(M.round, self, decimals)
 
+    @derived_from(pd.DataFrame)
+    def nunique(self, split_every=False, dropna=True, axis=0):
+        if axis == 1:
+            # split_every not used for axis=1
+            meta = self._meta_nonempty.nunique(axis=axis)
+            return self.map_partitions(
+                M.nunique,
+                meta=meta,
+                token="series-nunique",
+                axis=axis,
+                dropna=dropna,
+                enforce_metadata=False,
+            )
+        else:
+            nunique_list = [
+                self[col].nunique(split_every=split_every, dropna=dropna)
+                for col in self.columns
+            ]
+            name = "series-" + tokenize(*nunique_list)
+            dsk = {
+                (name, 0): (
+                    apply,
+                    pd.Series,
+                    [[(s._name, 0) for s in nunique_list]],
+                    {"index": self.columns},
+                )
+            }
+            graph = HighLevelGraph.from_collections(
+                name, dsk, dependencies=nunique_list
+            )
+            return Series(graph, name, self._meta.nunique(), (None, None))
+
     @derived_from(pd.DataFrame)
     def mode(self, dropna=True, split_every=False):
         mode_series_list = []


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
@@ -4752,6 +4752,19 @@ def test_pop():
     assert_eq(ddf, df[["x"]])
 
 
+@pytest.mark.parametrize("dropna", [True, False])
+@pytest.mark.parametrize("axis", [0, 1])
+def test_nunique(dropna, axis):
+    df = pd.DataFrame(
+        {"x": ["a", "a", "c"], "y": [None, 1, 2], "c": np.arange(0, 1, 0.4)}
+    )
+    ddf = dd.from_pandas(df, npartitions=2)
+    assert_eq(ddf["y"].nunique(dropna=dropna), df["y"].nunique(dropna=dropna))
+    assert_eq(
+        ddf.nunique(dropna=dropna, axis=axis), df.nunique(dropna=dropna, axis=axis)
+    )
+
+
 def test_simple_map_partitions():
     data = {"col_0": [9, -3, 0, -1, 5], "col_1": [-2, -7, 6, 8, -5]}
     df = pd.DataFrame(data)