Fix nunique for MultiIndex, DataFrame, and all NA case with dropna=False

cpp/src/stream_compaction/distinct_count.cu

@@ -187,7 +187,15 @@ cudf::size_type distinct_count(column_view const& input,
                                nan_policy nan_handling,
                                rmm::cuda_stream_view stream)
 {
-  if (0 == input.size() or input.null_count() == input.size()) { return 0; }
+  if (0 == input.size()) { return 0; }
+
+  if (input.null_count() == input.size()) {
+    if (null_handling == null_policy::INCLUDE) {
+      return 1;
+    } else {
+      return 0;
+    }
+  }
 
   auto count = detail::distinct_count(table_view{{input}}, null_equality::EQUAL, stream);
 

python/cudf/cudf/core/dataframe.py

@@ -7475,7 +7475,7 @@ def __dataframe__(
             self, nan_as_null=nan_as_null, allow_copy=allow_copy
         )
 
-    def nunique(self, axis=0, dropna=True):
+    def nunique(self, axis=0, dropna: bool = True) -> Series:
         """
         Count number of distinct elements in specified axis.
         Return Series with number of distinct elements. Can ignore NaN values.
@@ -7503,8 +7503,10 @@ def nunique(self, axis=0, dropna=True):
         """
         if axis != 0:
             raise NotImplementedError("axis parameter is not supported yet.")
-
-        return cudf.Series(super().nunique(dropna=dropna))
+        counts = [col.distinct_count(dropna=dropna) for col in self._columns]
+        return self._constructor_sliced(
+            counts, index=self._data.to_pandas_index()
+        )
 
     def _sample_axis_1(
         self,

python/cudf/cudf/core/frame.py

@@ -1940,10 +1940,9 @@ def nunique(self, dropna: bool = True):
         dict
             Name and unique value counts of each column in frame.
         """
-        return {
-            name: col.distinct_count(dropna=dropna)
-            for name, col in self._data.items()
-        }
+        raise NotImplementedError(
+            f"{type(self).__name__} does not implement nunique"
+        )
 
     @staticmethod
     @_cudf_nvtx_annotate

python/cudf/cudf/core/index.py

@@ -895,7 +895,7 @@ def __array__(self, dtype=None):
         )
 
     @_cudf_nvtx_annotate
-    def nunique(self) -> int:
+    def nunique(self, dropna: bool = True) -> int:
         return len(self)
 
     @_cudf_nvtx_annotate

python/cudf/cudf/core/multiindex.py

@@ -1701,6 +1701,14 @@ def fillna(self, value):
     def unique(self):
         return self.drop_duplicates(keep="first")
 
+    @_cudf_nvtx_annotate
+    def nunique(self, dropna: bool = True) -> int:
+        if dropna:
+            mi = self.dropna(how="all")
+        else:
+            mi = self
+        return len(mi.unique())
+
     def _clean_nulls_from_index(self):
         """
         Convert all na values(if any) in MultiIndex object

python/cudf/cudf/core/single_column_frame.py

@@ -335,8 +335,6 @@ def nunique(self, dropna: bool = True) -> int:
         int
             Number of unique values in the column.
         """
-        if self._column.null_count == len(self):
-            return 0
         return self._column.distinct_count(dropna=dropna)
 
     def _get_elements_from_column(self, arg) -> Union[ScalarLike, ColumnBase]:

python/cudf/cudf/tests/test_dataframe.py

@@ -9966,6 +9966,20 @@ def test_dataframe_nunique(data):
     assert_eq(expected, actual)
 
 
+@pytest.mark.parametrize(
+    "columns",
+    [
+        pd.RangeIndex(2, name="foo"),
+        pd.MultiIndex.from_arrays([[1, 2], [2, 3]], names=["foo", 1]),
+        pd.Index([3, 5], dtype=np.int8, name="foo"),
+    ],
+)
+def test_nunique_preserve_column_in_index(columns):
+    df = cudf.DataFrame([[1, 2]], columns=columns)
+    result = df.nunique().index.to_pandas()
+    pd.testing.assert_index_equal(result, columns, exact=True)
+
+
 @pytest.mark.parametrize(
     "data",
     [{"key": [0, 1, 1, 0, 0, 1], "val": [1, 8, 3, 9, -3, 8]}],

python/cudf/cudf/tests/test_multiindex.py

@@ -2162,3 +2162,14 @@ def test_multi_index_contains_hashable():
         lfunc_args_and_kwargs=((),),
         rfunc_args_and_kwargs=((),),
     )
+
+
+@pytest.mark.parametrize("array", [[1, 2], [1, np.nan], [np.nan] * 2])
+@pytest.mark.parametrize("dropna", [True, False])
+def test_nunique(array, dropna):
+    arrays = [array, [3, 4]]
+    gidx = cudf.MultiIndex.from_arrays(arrays)
+    pidx = pd.MultiIndex.from_arrays(arrays)
+    result = gidx.nunique(dropna=dropna)
+    expected = pidx.nunique(dropna=dropna)
+    assert result == expected

python/cudf/cudf/tests/test_series.py

@@ -2848,3 +2848,13 @@ def test_nans_to_nulls_noop_copies_column(value):
     ser1 = cudf.Series([value])
     ser2 = ser1.nans_to_nulls()
     assert ser1._column is not ser2._column
+
+
+@pytest.mark.parametrize("dropna", [False, True])
+def test_nunique_all_null(dropna):
+    data = [np.nan, np.nan]
+    pd_ser = pd.Series(data)
+    cudf_ser = cudf.Series(data, nan_as_null=True)
+    result = pd_ser.nunique(dropna=dropna)
+    expected = cudf_ser.nunique(dropna=dropna)
+    assert result == expected