API: add "level=" argument to MultiIndex.unique()

closes pandas-dev#17896

Author: toobaz

doc/source/whatsnew/v0.22.0.txt

@@ -102,7 +102,7 @@ Indexing
 ^^^^^^^^
 
 - Bug in :func:`PeriodIndex.truncate` which raises ``TypeError`` when ``PeriodIndex`` is monotonic (:issue:`17717`)
--
+- :func:`MultiIndex.unique` now supports the ``level=`` argument (:issue:`17896`)
 -
 
 I/O

pandas/core/indexes/base.py

@@ -3758,7 +3758,9 @@ def drop(self, labels, errors='raise'):
         return self.delete(indexer)
 
     @Appender(base._shared_docs['unique'] % _index_doc_kwargs)
-    def unique(self):
+    def unique(self, level=None):
+        if level not in {0, self.name, None}:
+            raise ValueError("Level {} not found".format(level))
         result = super(Index, self).unique()
         return self._shallow_copy(result)
 

pandas/core/indexes/multi.py

@@ -886,7 +886,7 @@ def _try_mi(k):
 
             raise InvalidIndexError(key)
 
-    def _get_level_values(self, level):
+    def _get_level_values(self, level, unique=False):
         """
         Return vector of label values for requested level,
         equal to the length of the index
@@ -896,17 +896,23 @@ def _get_level_values(self, level):
         Parameters
         ----------
         level : int level
+        unique : bool
+            if True, drop duplicated values
+
+            .. versionadded:: 0.22.0
 
         Returns
         -------
         values : ndarray
         """
 
-        unique = self.levels[level]
+        values = self.levels[level]
         labels = self.labels[level]
-        filled = algos.take_1d(unique._values, labels,
-                               fill_value=unique._na_value)
-        values = unique._shallow_copy(filled)
+        if unique:
+            labels = algos.unique(labels)
+        filled = algos.take_1d(values._values, labels,
+                               fill_value=values._na_value)
+        values = values._shallow_copy(filled)
         return values
 
     def get_level_values(self, level):
@@ -945,6 +951,35 @@ def get_level_values(self, level):
         values = self._get_level_values(level)
         return values
 
+    def unique(self, level=None):
+        """
+        Return unique values in the index. Uniques are returned in order
+        of appearance, this does NOT sort.
+
+        Parameters
+        ----------
+        level : int or str, optional, defaults None
+            only return values from specified level
+
+            .. versionadded:: 0.21.0
+
+        Returns
+        -------
+        Index without duplicates (MultiIndex as long as level=None)
+
+        See Also
+        --------
+        unique
+        Index.unique
+        Series.unique
+        """
+
+        if level is None:
+            return super(MultiIndex, self).unique()
+        else:
+            level = self._get_level_number(level)
+            return self._get_level_values(level=level, unique=True)
+
     def format(self, space=2, sparsify=None, adjoin=True, names=False,
                na_rep=None, formatter=None):
         if len(self) == 0:

pandas/tests/indexes/common.py

@@ -329,6 +329,25 @@ def test_duplicates(self, indices):
         assert not idx.is_unique
         assert idx.has_duplicates
 
+    def test_unique(self):
+        # GH 17896
+        idx = pd.Index([2, 3, 2, 1], name='my_index')
+        expected = pd.Index([2, 3, 1], name='my_index')
+        for level in 0, 'my_index', None:
+            result = idx.unique(level=level)
+            tm.assert_index_equal(result, expected)
+
+        for level in 3, 'wrong':
+            msg = "Level {} not found".format(level)
+            with tm.assert_raises_regex(ValueError, msg):
+                idx.unique(level=level)
+
+        # with NaNs
+        idx = pd.Index([2, np.nan, 2, 1], name='my_index')
+        expected = pd.Index([2, np.nan, 1], name='my_index')
+        result = idx.unique()
+        tm.assert_index_equal(result, expected)
+
     def test_get_unique_index(self, indices):
         # MultiIndex tested separately
         if not len(indices) or isinstance(indices, MultiIndex):

pandas/tests/indexes/test_multi.py

@@ -955,19 +955,21 @@ def test_get_level_values(self):
         exp = CategoricalIndex([1, 2, 3, 1, 2, 3])
         tm.assert_index_equal(index.get_level_values(1), exp)
 
-    def test_get_level_values_na(self):
+    @pytest.mark.xfail(reason='GH 17924 (returns Int64Index with float data)')
+    def test_get_level_values_int_with_na(self):
         arrays = [['a', 'b', 'b'], [1, np.nan, 2]]
         index = pd.MultiIndex.from_arrays(arrays)
-        values = index.get_level_values(1)
-        expected = np.array([1, np.nan, 2])
-        tm.assert_numpy_array_equal(values.values.astype(float), expected)
+        result = index.get_level_values(1)
+        expected = Index([1, np.nan, 2])
+        tm.assert_index_equal(result, expected)
 
         arrays = [['a', 'b', 'b'], [np.nan, np.nan, 2]]
         index = pd.MultiIndex.from_arrays(arrays)
-        values = index.get_level_values(1)
-        expected = np.array([np.nan, np.nan, 2])
-        tm.assert_numpy_array_equal(values.values.astype(float), expected)
+        result = index.get_level_values(1)
+        expected = Index([np.nan, np.nan, 2])
+        tm.assert_index_equal(result, expected)
 
+    def test_get_level_values_na(self):
         arrays = [[np.nan, np.nan, np.nan], ['a', np.nan, 1]]
         index = pd.MultiIndex.from_arrays(arrays)
         result = index.get_level_values(0)
@@ -982,7 +984,7 @@ def test_get_level_values_na(self):
         index = pd.MultiIndex.from_arrays(arrays)
         values = index.get_level_values(1)
         expected = pd.DatetimeIndex([0, 1, pd.NaT])
-        tm.assert_numpy_array_equal(values.values, expected.values)
+        tm.assert_index_equal(values, expected)
 
         arrays = [[], []]
         index = pd.MultiIndex.from_arrays(arrays)
@@ -2269,6 +2271,28 @@ def test_unique(self):
         exp = pd.MultiIndex.from_arrays([['a'], ['a']])
         tm.assert_index_equal(res, exp)
 
+        # GH #17896 - with level= argument
+        for level in 0, 'first':
+            result = self.index.unique(level=level)
+            expected = Index(['foo', 'bar', 'baz', 'qux'],
+                             name='first')
+            tm.assert_index_equal(result, expected)
+
+            # With already unique level
+            mi = pd.MultiIndex.from_arrays([[1, 3, 2, 4], [1, 1, 1, 2]],
+                                           names=['first', 'second'])
+            result = mi.unique(level=level)
+            expected = mi.get_level_values(level)
+            tm.assert_index_equal(result, expected)
+
+    @pytest.mark.xfail(reason='GH 17924 (returns Int64Index with float data)')
+    def test_unique_with_nans(self):
+        # with NaNs
+        index = pd.MultiIndex.from_arrays([['a', 'b', 'b'], [2, np.nan, 2]])
+        result = index.unique(level=1)
+        expected = pd.Index([2, np.nan])
+        tm.assert_index_equal(result, expected)
+
     def test_unique_datetimelike(self):
         idx1 = pd.DatetimeIndex(['2015-01-01', '2015-01-01', '2015-01-01',
                                  '2015-01-01', 'NaT', 'NaT'])