Implement Index.nunique

databricks/koalas/indexes.py

@@ -588,6 +588,48 @@ def dropna(self):
         kdf = DataFrame(internal)
         return Index(kdf) if type(self) == Index else MultiIndex(kdf)
 
+    def nunique(self, dropna=True):
+        """
+        Return number of unique elements in the object.
+
+        Excludes NA values by default.
+
+        Parameters
+        ----------
+        dropna : bool, default True
+            Don't include NaN in the count.
+
+        Returns
+        -------
+        int
+
+        See Also
+        --------
+        DataFrame.nunique: Method nunique for DataFrame.
+        Series.count: Count non-NA/null observations in the Series.
+
+        Examples
+        --------
+        >>> idx = ks.Index([1, 1, 2, None])
+        >>> idx
+        Float64Index([1.0, 1.0, 2.0, nan], dtype='float64')
+
+        >>> idx.nunique()
+        2
+
+        >>> idx.nunique(dropna=False)
+        3
+        """
+        index_scol = self._internal.index_scols[0]
+        sdf = self._internal._sdf
+        if dropna:
+            result = F.countDistinct(index_scol)
+        else:
+            result = (F.countDistinct(index_scol) +
+                      F.when(F.count(F.when(index_scol.isNull(), 1)
+                                     .otherwise(None)) >= 1, 1).otherwise(0))
+        return sdf.select([result]).collect()[0][0]
+
     def unique(self, level=None):
         """
         Return unique values in the index.

databricks/koalas/missing/indexes.py

@@ -73,7 +73,6 @@ class _MissingPandasLikeIndex(object):
     is_type_compatible = unsupported_function('is_type_compatible')
     join = unsupported_function('join')
     map = unsupported_function('map')
-    nunique = unsupported_function('nunique')
     putmask = unsupported_function('putmask')
     ravel = unsupported_function('ravel')
     reindex = unsupported_function('reindex')

databricks/koalas/tests/test_indexes.py

@@ -357,3 +357,10 @@ def test_multiindex_isna(self):
                 NotImplementedError,
                 "notna is not defined for MultiIndex"):
             kidx.notnull()
+
+    def test_index_nunique(self):
+        pidx = pd.DataFrame({'a': ['a', 'b', 'c', 'd']}, index=[1, 1, 2, None]).index
+        kidx = ks.DataFrame({'a': ['a', 'b', 'c', 'd']}, index=[1, 1, 2, None]).index
+
+        self.assert_eq(pidx.nunique(), kidx.nunique())
+        self.assert_eq(pidx.nunique(dropna=True), kidx.nunique(dropna=True))

docs/source/reference/indexing.rst

@@ -53,6 +53,7 @@ Modifying and computations
    Index.max
    Index.rename
    Index.unique
+   Index.nunique
    Index.value_counts
 
 Missing Values