@@ -116,6 +116,25 @@ def ref_locs(self):
116
116
self ._ref_locs = indexer
117
117
return self ._ref_locs
118
118
119
+ def take_ref_locs (self , indexer ):
120
+ """
121
+ need to preserve the ref_locs and just shift them
122
+ return None if ref_locs is None
123
+
124
+ see GH6509
125
+ """
126
+
127
+ ref_locs = self ._ref_locs
128
+ if ref_locs is None :
129
+ return None
130
+
131
+ tindexer = np .ones (len (ref_locs ),dtype = bool )
132
+ tindexer [indexer ] = False
133
+ tindexer = tindexer .astype (int ).cumsum ()[indexer ]
134
+ ref_locs = ref_locs [indexer ]
135
+ ref_locs -= tindexer
136
+ return ref_locs
137
+
119
138
def reset_ref_locs (self ):
120
139
""" reset the block ref_locs """
121
140
self ._ref_locs = np .empty (len (self .items ), dtype = 'int64' )
@@ -866,13 +885,20 @@ def func(x):
866
885
ndim = self .ndim , klass = self .__class__ , fastpath = True )]
867
886
return self ._maybe_downcast (blocks , downcast )
868
887
869
- def take (self , indexer , ref_items , axis = 1 ):
888
+ def take (self , indexer , ref_items , new_axis , axis = 1 ):
870
889
if axis < 1 :
871
890
raise AssertionError ('axis must be at least 1, got %d' % axis )
872
891
new_values = com .take_nd (self .values , indexer , axis = axis ,
873
892
allow_fill = False )
893
+
894
+ # need to preserve the ref_locs and just shift them
895
+ # GH6121
896
+ ref_locs = None
897
+ if not new_axis .is_unique :
898
+ ref_locs = self ._ref_locs
899
+
874
900
return [make_block (new_values , self .items , ref_items , ndim = self .ndim ,
875
- klass = self .__class__ , fastpath = True )]
901
+ klass = self .__class__ , placement = ref_locs , fastpath = True )]
876
902
877
903
def get_values (self , dtype = None ):
878
904
return self .values
@@ -1820,7 +1846,7 @@ def shift(self, indexer, periods, axis=0):
1820
1846
new_values [periods :] = fill_value
1821
1847
return [self .make_block (new_values )]
1822
1848
1823
- def take (self , indexer , ref_items , axis = 1 ):
1849
+ def take (self , indexer , ref_items , new_axis , axis = 1 ):
1824
1850
""" going to take our items
1825
1851
along the long dimension"""
1826
1852
if axis < 1 :
@@ -2601,18 +2627,7 @@ def get_slice(self, slobj, axis=0, raise_on_error=False):
2601
2627
if len (self .blocks ) == 1 :
2602
2628
2603
2629
blk = self .blocks [0 ]
2604
-
2605
- # see GH 6059
2606
- ref_locs = blk ._ref_locs
2607
- if ref_locs is not None :
2608
-
2609
- # need to preserve the ref_locs and just shift them
2610
- indexer = np .ones (len (ref_locs ),dtype = bool )
2611
- indexer [slobj ] = False
2612
- indexer = indexer .astype (int ).cumsum ()[slobj ]
2613
- ref_locs = ref_locs [slobj ]
2614
- ref_locs -= indexer
2615
-
2630
+ ref_locs = blk .take_ref_locs (slobj )
2616
2631
newb = make_block (blk ._slice (slobj ), new_items , new_items ,
2617
2632
klass = blk .__class__ , fastpath = True ,
2618
2633
placement = ref_locs )
@@ -3371,6 +3386,7 @@ def take(self, indexer, new_index=None, axis=1, verify=True):
3371
3386
if axis < 1 :
3372
3387
raise AssertionError ('axis must be at least 1, got %d' % axis )
3373
3388
3389
+ self ._consolidate_inplace ()
3374
3390
if isinstance (indexer , list ):
3375
3391
indexer = np .array (indexer )
3376
3392
@@ -3388,8 +3404,12 @@ def take(self, indexer, new_index=None, axis=1, verify=True):
3388
3404
new_index = self .axes [axis ].take (indexer )
3389
3405
3390
3406
new_axes [axis ] = new_index
3391
- return self .apply ('take' , axes = new_axes , indexer = indexer ,
3392
- ref_items = new_axes [0 ], axis = axis )
3407
+ return self .apply ('take' ,
3408
+ axes = new_axes ,
3409
+ indexer = indexer ,
3410
+ ref_items = new_axes [0 ],
3411
+ new_axis = new_axes [axis ],
3412
+ axis = axis )
3393
3413
3394
3414
def merge (self , other , lsuffix = None , rsuffix = None ):
3395
3415
if not self ._is_indexed_like (other ):
0 commit comments