Implement MultiIndex.levshape (#1086)

Get MultiIndex unique value count in every level.

Author: RainFung

databricks/koalas/indexes.py

@@ -861,6 +861,27 @@ def name(self) -> str:
     def name(self, name: str) -> None:
         raise PandasNotImplementedError(class_name='pd.MultiIndex', property_name='name')
 
+    @property
+    def levshape(self):
+        """
+        A tuple with the length of each level.
+
+        Examples
+        --------
+        >>> midx = ks.MultiIndex.from_tuples([('a', 'x'), ('b', 'y'), ('c', 'z')])
+        >>> midx  # doctest: +SKIP
+        MultiIndex([('a', 'x'),
+                    ('b', 'y'),
+                    ('c', 'z')],
+                   )
+
+        >>> midx.levshape
+        (3, 3)
+        """
+        internal = self._internal
+        result = internal._sdf.agg(*(F.countDistinct(c) for c in internal.index_scols)).collect()[0]
+        return tuple(result)
+
     def to_pandas(self) -> pd.MultiIndex:
         """
         Return a pandas MultiIndex.

databricks/koalas/missing/indexes.py

@@ -118,7 +118,6 @@ class _MissingPandasLikeMultiIndex(object):
 
     # Properties
     is_all_dates = unsupported_property('is_all_dates')
-    levshape = unsupported_property('levshape')
 
     # Deprecated properties
     strides = unsupported_property('strides', deprecated=True)

databricks/koalas/tests/test_indexes.py

@@ -150,6 +150,11 @@ def test_multi_index_names(self):
         with self.assertRaises(PandasNotImplementedError):
             kidx.name = 'renamed'
 
+    def test_multi_index_levshape(self):
+        pidx = pd.MultiIndex.from_tuples([('a', 'x', 1), ('b', 'y', 2)])
+        kidx = ks.MultiIndex.from_tuples([('a', 'x', 1), ('b', 'y', 2)])
+        self.assertEqual(pidx.levshape, kidx.levshape)
+
     def test_index_unique(self):
         kidx = self.kdf.index
 

docs/source/reference/indexing.rst

@@ -101,6 +101,7 @@ MultiIndex Properties
    MultiIndex.ndim
    MultiIndex.T
    MultiIndex.nlevels
+   MultiIndex.levshape
 
 MultiIndex Missing Values
 ~~~~~~~~~~~~~~~~~~~~~~~~~