Refactor the code to work with period time time delta indices

Author: nateyoder

doc/source/whatsnew/v0.20.0.txt

@@ -139,6 +139,8 @@ Other enhancements
 - ``pd.merge_asof()`` gained the option ``direction='backward'|'forward'|'nearest'`` (:issue:`14887`)
 - ``Series/DataFrame.asfreq()`` have gained a ``fill_value`` parameter, to fill missing values (:issue:`3715`).
 - ``Series/DataFrame.resample.asfreq`` have gained a ``fill_value`` parameter, to fill missing values during resampling (:issue:`3715`).
+- ``Index.map`` can now accept series and dictionary input object.
+
 
 .. _ISO 8601 duration: https://en.wikipedia.org/wiki/ISO_8601#Durations
 

pandas/core/series.py

@@ -184,23 +184,9 @@ def __init__(self, data=None, index=None, dtype=None, name=None,
                         index = Index(data)
                     else:
                         index = Index(_try_sort(data))
+
                 try:
-                    if isinstance(index, DatetimeIndex):
-                        if len(data):
-                            # coerce back to datetime objects for lookup
-                            data = _dict_compat(data)
-                            data = lib.fast_multiget(data,
-                                                     index.asobject.values,
-                                                     default=np.nan)
-                        else:
-                            data = np.nan
-                    # GH #12169
-                    elif isinstance(index, (PeriodIndex, TimedeltaIndex)):
-                        data = ([data.get(i, nan) for i in index]
-                                if data else np.nan)
-                    else:
-                        data = lib.fast_multiget(data, index.values,
-                                                 default=np.nan)
+                    data = index.get_values_from_dict(data)
                 except TypeError:
                     data = ([data.get(i, nan) for i in index]
                             if data else np.nan)

pandas/indexes/base.py

@@ -2448,6 +2448,25 @@ def get_indexer_for(self, target, **kwargs):
         indexer, _ = self.get_indexer_non_unique(target, **kwargs)
         return indexer
 
+    def get_values_from_dict(self, input_dict):
+        """Return the values of the input dictionary in the order the keys are
+        in the index. np.nan is returned for index values not in the
+        dictionary.
+
+        Parameters
+        ----------
+        input_dict : dict
+            The dictionary from which to extract the values
+
+        Returns
+        -------
+        Union[np.array, list]
+
+        """
+
+        return lib.fast_multiget(input_dict, self.values,
+                                 default=np.nan)
+
     def _possibly_promote(self, other):
         # A hack, but it works
         from pandas.tseries.index import DatetimeIndex
@@ -2491,8 +2510,8 @@ def map(self, mapper):
 
         Parameters
         ----------
-        mapper : function, dict, or Series
-            Function to be applied.
+        mapper : Union[function, dict, Series]
+            Function to be applied or input correspondence object.
 
         Returns
         -------
@@ -2505,12 +2524,15 @@ def map(self, mapper):
         from .multi import MultiIndex
 
         if isinstance(mapper, ABCSeries):
-            indexer = mapper.index.get_indexer(self._values)
+            indexer = mapper.index.get_indexer(self.values)
             mapped_values = algos.take_1d(mapper.values, indexer)
+        elif isinstance(mapper, dict):
+            idx = Index(mapper.keys())
+            data = idx.get_values_from_dict(mapper)
+            indexer = idx.get_indexer(self.values)
+            mapped_values = algos.take_1d(data, indexer)
         else:
-            if isinstance(mapper, dict):
-                mapper = mapper.get
-            mapped_values = self._arrmap(self._values, mapper)
+            mapped_values = self._arrmap(self.values, mapper)
 
         attributes = self._get_attributes_dict()
         if mapped_values.size and isinstance(mapped_values[0], tuple):

pandas/tests/indexes/test_base.py

@@ -811,10 +811,10 @@ 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):
+    def test_map_with_dict_and_series(self):
         expected = Index(['foo', 'bar', 'baz'])
         mapper = Series(expected.values, index=[0, 1, 2])
-        self.assert_index_equal(tm.makeIntIndex(3).map(mapper), expected)
+        tm.assert_index_equal(tm.makeIntIndex(3).map(mapper), expected)
 
         # GH 12766
         # special = []
@@ -824,41 +824,58 @@ def test_map_with_series_all_indices(self):
             orig_values = ['a', 'B', 1, 'a']
             new_values = ['one', 2, 3.0, 'one']
             cur_index = CategoricalIndex(orig_values, name='XXX')
+            expected = CategoricalIndex(new_values,
+                                        name='XXX', categories=[3.0, 2, 'one'])
+
             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)
+            # Order of categories in output can be different
+            tm.assert_index_equal(expected, output)
+
+            mapper = {o: n for o, n in
+                      zip(orig_values[:-1], new_values[:-1])}
+            output = cur_index.map(mapper)
+            # Order of categories in output can be different
+            tm.assert_index_equal(expected, output)
 
         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))
+            mapper = pd.Series(expected, index=cur_index)
+            tm.assert_index_equal(expected, cur_index.map(mapper))
+
+            mapper = {o: n for o, n in
+                      zip(cur_index, expected)}
+            if mapper:
+                tm.assert_index_equal(expected, cur_index.map(mapper))
+            else:
+                # The expected index type is Int64Index
+                # but the output defaults to Float64
+                tm.assert_index_equal(Float64Index([]),
+                                      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")
+        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)
+        tm.assert_index_equal(a.map(b), exp)
         exp = Index(["odd", "even", "odd", np.nan])
-        self.assert_index_equal(a.map(c), exp)
+        tm.assert_index_equal(a.map(c), exp)
 
     def test_map_with_non_function_missing_values(self):
         # GH 12756
         expected = Index([2., np.nan, 'foo'])
         input = Index([2, 1, 0])
 
         mapper = Series(['foo', 2., 'baz'], index=[0, 2, -1])
-        self.assert_index_equal(expected, input.map(mapper))
+        tm.assert_index_equal(expected, input.map(mapper))
 
         mapper = {0: 'foo', 2: 2.0, -1: 'baz'}
-        self.assert_index_equal(expected, input.map(mapper))
+        tm.assert_index_equal(expected, input.map(mapper))
 
     def test_append_multiple(self):
         index = Index(['a', 'b', 'c', 'd', 'e', 'f'])

pandas/tseries/index.py

@@ -1380,6 +1380,30 @@ def get_value_maybe_box(self, series, key):
                                         key, tz=self.tz)
         return _maybe_box(self, values, series, key)
 
+    def get_values_from_dict(self, input_dict):
+        """Return the values of the input dictionary in the order the keys are
+        in the index. np.nan is returned for index values not in the
+        dictionary.
+
+        Parameters
+        ----------
+        input_dict : dict
+            The dictionary from which to extract the values
+
+        Returns
+        -------
+        Union[np.array, list]
+
+        """
+        if len(input_dict):
+            # coerce back to datetime objects for lookup
+            input_dict = com._dict_compat(input_dict)
+            return lib.fast_multiget(input_dict,
+                                     self.asobject.values,
+                                     default=np.nan)
+        else:
+            return np.nan
+
     def get_loc(self, key, method=None, tolerance=None):
         """
         Get integer location for requested label

pandas/tseries/period.py

@@ -783,6 +783,25 @@ def _get_unique_index(self, dropna=False):
             res = res.dropna()
         return res
 
+    def get_values_from_dict(self, input_dict):
+        """Return the values of the input dictionary in the order the keys are
+        in the index. np.nan is returned for index values not in the
+        dictionary.
+
+        Parameters
+        ----------
+        input_dict : dict
+            The dictionary from which to extract the values
+
+        Returns
+        -------
+        Union[np.array, list]
+
+        """
+
+        return np.array([input_dict.get(i, np.nan) for i in self.values]
+                        if input_dict else [np.nan])
+
     def get_loc(self, key, method=None, tolerance=None):
         """
         Get integer location for requested label

pandas/tseries/tdi.py

@@ -665,6 +665,26 @@ def get_value_maybe_box(self, series, key):
         values = self._engine.get_value(_values_from_object(series), key)
         return _maybe_box(self, values, series, key)
 
+    def get_values_from_dict(self, input_dict):
+        """Return the values of the input dictionary in the order the keys are
+        in the index. np.nan is returned for index values not in the
+        dictionary.
+
+        Parameters
+        ----------
+        input_dict : dict
+            The dictionary from which to extract the values
+
+        Returns
+        -------
+        Union[np.array, list]
+
+        """
+
+        return np.array([input_dict.get(i, np.nan)
+                         for i in self.asobject.values]
+                        if input_dict else [np.nan])
+
     def get_loc(self, key, method=None, tolerance=None):
         """
         Get integer location for requested label

pandas/tseries/tests/test_period.py

@@ -3511,6 +3511,10 @@ def test_map(self):
         result = index.map(lambda x: x + 1)
         expected = index + 1
         tm.assert_index_equal(result, expected)
+        series_map = pd.Series(expected, index)
+        tm.assert_index_equal(index.map(series_map), expected)
+        dict_map = {i: e for e, i in zip(expected, index)}
+        tm.assert_index_equal(index.map(dict_map), expected)
 
         result = index.map(lambda x: x.ordinal)
         exp = Index([x.ordinal for x in index])

pandas/tseries/tests/test_timedeltas.py

@@ -1574,14 +1574,19 @@ def test_ops_error_str(self):
                 l != r
 
     def test_map(self):
-
         rng = timedelta_range('1 day', periods=10)
 
         f = lambda x: x.days
         result = rng.map(f)
         exp = Int64Index([f(x) for x in rng])
         tm.assert_index_equal(result, exp)
 
+        s = Series(exp, index=rng)
+        tm.assert_index_equal(rng.map(s), exp)
+
+        d = {r: v for r, v in zip(rng, exp)}
+        tm.assert_index_equal(rng.map(d), exp)
+
     def test_misc_coverage(self):
 
         rng = timedelta_range('1 day', periods=5)

pandas/tseries/tests/test_timeseries.py

@@ -3012,9 +3012,16 @@ def test_map(self):
         rng = date_range('1/1/2000', periods=10)
 
         f = lambda x: x.strftime('%Y%m%d')
-        result = rng.map(f)
-        exp = Index([f(x) for x in rng], dtype='<U8')
-        tm.assert_index_equal(result, exp)
+        exp_values = [f(x) for x in rng]
+        exp = Index(exp_values, dtype='<U8')
+
+        tm.assert_index_equal(rng.map(f), exp)
+
+        mapper = {o: n for o, n in zip(rng, exp_values)}
+        tm.assert_index_equal(rng.map(mapper), exp)
+
+        mapper = Series(exp_values, index=rng)
+        tm.assert_index_equal(rng.map(mapper), exp)
 
     def test_iteration_preserves_tz(self):