DataFrame.truncate drops MultiIndex names (#34589)

Author: arw2019

doc/source/whatsnew/v1.1.0.rst

@@ -956,6 +956,7 @@ MultiIndex
  df.loc[(['b', 'a'], [2, 1]), :]
 
 - Bug in :meth:`MultiIndex.intersection` was not guaranteed to preserve order when ``sort=False``. (:issue:`31325`)
+- Bug in :meth:`DataFrame.truncate` was dropping :class:`MultiIndex` names. (:issue:`34564`)
 
 .. ipython:: python
 

pandas/core/indexes/multi.py

@@ -3193,7 +3193,12 @@ def truncate(self, before=None, after=None):
  new_codes = [level_codes[left:right] for level_codes in self.codes]
  new_codes[0] = new_codes[0] - i
 
- return MultiIndex(levels=new_levels, codes=new_codes, verify_integrity=False)
+ return MultiIndex(
+ levels=new_levels,
+ codes=new_codes,
+ names=self._names,
+ verify_integrity=False,
+ )
 
  def equals(self, other) -> bool:
  """

pandas/tests/frame/methods/test_truncate.py

@@ -104,3 +104,16 @@ def test_truncate_decreasing_index(self, before, after, indices, klass):
  result = values.truncate(before=before, after=after)
  expected = values.loc[indices]
  tm.assert_frame_equal(result, expected)
+
+ def test_truncate_multiindex(self):
+ # GH 34564
+ mi = pd.MultiIndex.from_product([[1, 2, 3, 4], ["A", "B"]], names=["L1", "L2"])
+ s1 = pd.DataFrame(range(mi.shape[0]), index=mi, columns=["col"])
+ result = s1.truncate(before=2, after=3)
+
+ df = pd.DataFrame.from_dict(
+ {"L1": [2, 2, 3, 3], "L2": ["A", "B", "A", "B"], "col": [2, 3, 4, 5]}
+ )
+ expected = df.set_index(["L1", "L2"])
+
+ tm.assert_frame_equal(result, expected)

pandas/tests/indexes/multi/test_analytics.py

@@ -30,27 +30,33 @@ def test_groupby(idx):
  tm.assert_dict_equal(groups, exp)
 
 
-def test_truncate():
+def test_truncate_multiindex():
+ # GH 34564 for MultiIndex level names check
  major_axis = Index(list(range(4)))
  minor_axis = Index(list(range(2)))
 
  major_codes = np.array([0, 0, 1, 2, 3, 3])
  minor_codes = np.array([0, 1, 0, 1, 0, 1])
 
  index = MultiIndex(
- levels=[major_axis, minor_axis], codes=[major_codes, minor_codes]
+ levels=[major_axis, minor_axis],
+ codes=[major_codes, minor_codes],
+ names=["L1", "L2"],
  )
 
  result = index.truncate(before=1)
  assert "foo" not in result.levels[0]
  assert 1 in result.levels[0]
+ assert index.names == result.names
 
  result = index.truncate(after=1)
  assert 2 not in result.levels[0]
  assert 1 in result.levels[0]
+ assert index.names == result.names
 
  result = index.truncate(before=1, after=2)
  assert len(result.levels[0]) == 2
+ assert index.names == result.names
 
  msg = "after < before"
  with pytest.raises(ValueError, match=msg):

pandas/tests/series/methods/test_truncate.py

@@ -126,3 +126,17 @@ def test_truncate_periodindex(self):
 
  expected_idx2 = pd.PeriodIndex([pd.Period("2017-09-02")])
  tm.assert_series_equal(result2, pd.Series([2], index=expected_idx2))
+
+ def test_truncate_multiindex(self):
+ # GH 34564
+ mi = pd.MultiIndex.from_product([[1, 2, 3, 4], ["A", "B"]], names=["L1", "L2"])
+ s1 = pd.Series(range(mi.shape[0]), index=mi, name="col")
+ result = s1.truncate(before=2, after=3)
+
+ df = pd.DataFrame.from_dict(
+ {"L1": [2, 2, 3, 3], "L2": ["A", "B", "A", "B"], "col": [2, 3, 4, 5]}
+ )
+ df.set_index(["L1", "L2"], inplace=True)
+ expected = df.col
+
+ tm.assert_series_equal(result, expected)