Skip to content

Commit a596744

Browse files
committed
introspect results from map so that if the output array has tuples we create a multiindex instead of an index
1 parent 5fc66c3 commit a596744

File tree

5 files changed

+41
-9
lines changed

5 files changed

+41
-9
lines changed

doc/source/whatsnew/v0.20.0.txt

+19
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,25 @@ Other enhancements
9191
Backwards incompatible API changes
9292
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
9393

94+
- ``map`` on an ``Index`` now returns an ``Index``, not an array (:issue:`12766`)
95+
.. ipython:: python
96+
97+
idx = Index([1, 2])
98+
idx
99+
100+
Previous Behavior:
101+
102+
.. code-block:: ipython
103+
104+
In [3]: idx.map(lambda x: x * 2)
105+
Out[3]: array([2, 4])
106+
107+
New Behavior:
108+
109+
.. ipython:: python
110+
111+
idx.map(lambda x: x * 2)
112+
94113
.. _whatsnew_0200.api:
95114

96115

pandas/indexes/base.py

+8-1
Original file line numberDiff line numberDiff line change
@@ -2439,10 +2439,17 @@ def map(self, mapper):
24392439
-------
24402440
applied : Index
24412441
The output of the mapping function applied to the index.
2442+
If the function returns a tuple a
24422443
"""
2444+
from .multi import MultiIndex
2445+
mapped_values = self._arrmap(self.values, mapper)
24432446
attributes = self._get_attributes_dict()
2447+
if mapped_values.size and isinstance(mapped_values[0], tuple):
2448+
return MultiIndex.from_tuples(mapped_values,
2449+
names=attributes.get('name'))
2450+
24442451
attributes['copy'] = False
2445-
return Index(self._arrmap(self.values, mapper), **attributes)
2452+
return Index(mapped_values, **attributes)
24462453

24472454
def isin(self, values, level=None):
24482455
"""

pandas/tests/indexes/test_base.py

+12-8
Original file line numberDiff line numberDiff line change
@@ -768,20 +768,24 @@ def test_sub(self):
768768
self.assertRaises(TypeError, lambda: idx.tolist() - idx)
769769

770770
def test_map_identity_mapping(self):
771+
# GH 12766
771772
for name, cur_index in self.indices.items():
772-
if name == 'tuples':
773-
expected = Index(cur_index.values, tupleize_cols=False)
774-
self.assert_index_equal(expected, cur_index.map(lambda x: x))
775-
else:
776-
self.assert_index_equal(cur_index, cur_index.map(lambda x: x))
773+
self.assert_index_equal(cur_index, cur_index.map(lambda x: x))
777774

778-
def test_map_that_returns_tuples_creates_index_not_multi_index(self):
775+
def test_map_that_returns_tuples_creates_multi_index(self):
776+
# GH 12766
779777
boolean_index = tm.makeIntIndex(3).map(lambda x: (x, x == 1))
780-
expected = Index([(0, False), (1, True), (2, False)],
781-
tupleize_cols=False)
778+
expected = MultiIndex.from_tuples([(0, False), (1, True), (2, False)])
779+
self.assert_index_equal(boolean_index, expected)
780+
781+
def test_map_that_returns_a_length_one_tuple_creates_an_index(self):
782+
# GH 12766
783+
boolean_index = tm.makeIntIndex(3).map(lambda x: (x, ))
784+
expected = Index([(0, ), (1, ), (2, )])
782785
self.assert_index_equal(boolean_index, expected)
783786

784787
def test_map_that_reduces_multi_index_to_single_index_returns_index(self):
788+
# GH 12766
785789
first_level = ['foo', 'bar', 'baz']
786790
multi_index = MultiIndex.from_tuples(lzip(first_level, [1, 2, 3]))
787791
reduced_index = multi_index.map(lambda x: x[0])

pandas/tests/indexes/test_category.py

+1
Original file line numberDiff line numberDiff line change
@@ -218,6 +218,7 @@ def test_map(self):
218218
ordered=False, name='XXX')
219219
tm.assert_index_equal(result, exp)
220220

221+
# GH 12766: Return an index not an array
221222
tm.assert_index_equal(ci.map(lambda x: 1),
222223
Index(np.array([1] * 5, dtype=np.int64), name='XXX'))
223224

pandas/tests/test_categorical.py

+1
Original file line numberDiff line numberDiff line change
@@ -1669,6 +1669,7 @@ def test_map(self):
16691669
tm.assert_categorical_equal(result, exp)
16701670

16711671
result = c.map(lambda x: 1)
1672+
# GH 12766: Return an index not an array
16721673
tm.assert_index_equal(result, Index(np.array([1] * 5, dtype=np.int64)))
16731674

16741675

0 commit comments

Comments
 (0)