Allow index.map() to accept series and dictionary inputs in addition to functional inputs

Author: nateyoder

pandas/indexes/base.py

@@ -14,7 +14,6 @@
 from pandas.compat.numpy import function as nv
 from pandas import compat
 
-
 from pandas.types.generic import ABCSeries, ABCMultiIndex, ABCPeriodIndex
 from pandas.types.missing import isnull, array_equivalent
 from pandas.types.common import (_ensure_int64,
@@ -2492,7 +2491,7 @@ def map(self, mapper):
 
         Parameters
         ----------
-        mapper : callable
+        mapper : function, dict, or Series
             Function to be applied.
 
         Returns
@@ -2504,7 +2503,15 @@ def map(self, mapper):
 
         """
         from .multi import MultiIndex
-        mapped_values = self._arrmap(self.values, mapper)
+
+        if isinstance(mapper, ABCSeries):
+            indexer = mapper.index.get_indexer(self._values)
+            mapped_values = algos.take_1d(mapper.values, indexer)
+        else:
+            if isinstance(mapper, dict):
+                mapper = mapper.get
+            mapped_values = self._arrmap(self._values, mapper)
+
         attributes = self._get_attributes_dict()
         if mapped_values.size and isinstance(mapped_values[0], tuple):
             return MultiIndex.from_tuples(mapped_values,

pandas/tests/indexes/test_base.py

@@ -811,6 +811,52 @@ def test_map_tseries_indices_return_index(self):
         exp = Index(range(24), name='hourly')
         tm.assert_index_equal(exp, date_index.map(lambda x: x.hour))
 
+    def test_map_with_series_all_indices(self):
+        expected = Index(['foo', 'bar', 'baz'])
+        mapper = Series(expected.values, index=[0, 1, 2])
+        self.assert_index_equal(tm.makeIntIndex(3).map(mapper), expected)
+
+        # GH 12766
+        # special = []
+        special = ['catIndex']
+
+        for name in special:
+            orig_values = ['a', 'B', 1, 'a']
+            new_values = ['one', 2, 3.0, 'one']
+            cur_index = CategoricalIndex(orig_values, name='XXX')
+            mapper = pd.Series(new_values[:-1], index=orig_values[:-1])
+            expected = CategoricalIndex(new_values, name='XXX')
+            output = cur_index.map(mapper)
+            self.assert_numpy_array_equal(expected.values.get_values(), output.values.get_values())
+            self.assert_equal(expected.name, output.name)
+
+
+        for name in list(set(self.indices.keys()) - set(special)):
+            cur_index = self.indices[name]
+            expected = Index(np.arange(len(cur_index), 0, -1))
+            mapper = pd.Series(expected.values, index=cur_index)
+            print(name)
+            output = cur_index.map(mapper)
+            self.assert_index_equal(expected, cur_index.map(mapper))
+
+    def test_map_with_categorical_series(self):
+        # GH 12756
+        a = Index([1, 2, 3, 4])
+        b = Series(["even", "odd", "even", "odd"], dtype="category")
+        c = Series(["even", "odd", "even", "odd"])
+
+        exp = CategoricalIndex(["odd", "even", "odd", np.nan])
+        self.assert_index_equal(a.map(b), exp)
+        exp = Index(["odd", "even", "odd", np.nan])
+        self.assert_index_equal(a.map(c), exp)
+
+    def test_map_with_series_missing_values(self):
+        # GH 12756
+        expected = Index([2., np.nan, 'foo'])
+        mapper = Series(['foo', 2., 'baz'], index=[0, 2, -1])
+        output = Index([2, 1, 0]).map(mapper)
+        self.assert_index_equal(output, expected)
+
     def test_append_multiple(self):
         index = Index(['a', 'b', 'c', 'd', 'e', 'f'])
 

pandas/tests/indexes/test_category.py

@@ -233,6 +233,12 @@ def f(x):
                                   ordered=False)
         tm.assert_index_equal(result, exp)
 
+        result = ci.map(pd.Series([10, 20, 30], index=['A', 'B', 'C']))
+        tm.assert_index_equal(result, exp)
+
+        result = ci.map({'A': 10, 'B': 20, 'C': 30})
+        tm.assert_index_equal(result, exp)
+
     def test_where(self):
         i = self.create_index()
         result = i.where(notnull(i))