@@ -1111,6 +1111,7 @@ def blk_func(block: "Block") -> List["Block"]:
1111
1111
assert len (locs ) == result .shape [1 ]
1112
1112
for i , loc in enumerate (locs ):
1113
1113
agg_block = result .iloc [:, [i ]]._mgr .blocks [0 ]
1114
+ agg_block .mgr_locs = [loc ]
1114
1115
new_blocks .append (agg_block )
1115
1116
else :
1116
1117
result = result ._mgr .blocks [0 ].values
@@ -1124,7 +1125,6 @@ def blk_func(block: "Block") -> List["Block"]:
1124
1125
return new_blocks
1125
1126
1126
1127
skipped : List [int ] = []
1127
- new_items : List [np .ndarray ] = []
1128
1128
for i , block in enumerate (data .blocks ):
1129
1129
try :
1130
1130
nbs = blk_func (block )
@@ -1136,33 +1136,13 @@ def blk_func(block: "Block") -> List["Block"]:
1136
1136
deleted_items .append (block .mgr_locs .as_array )
1137
1137
else :
1138
1138
agg_blocks .extend (nbs )
1139
- new_items .append (block .mgr_locs .as_array )
1140
1139
1141
1140
if not agg_blocks :
1142
1141
raise DataError ("No numeric types to aggregate" )
1143
1142
1144
1143
# reset the locs in the blocks to correspond to our
1145
1144
# current ordering
1146
- indexer = np .concatenate (new_items )
1147
- agg_items = data .items .take (np .sort (indexer ))
1148
-
1149
- if deleted_items :
1150
-
1151
- # we need to adjust the indexer to account for the
1152
- # items we have removed
1153
- # really should be done in internals :<
1154
-
1155
- deleted = np .concatenate (deleted_items )
1156
- ai = np .arange (len (data ))
1157
- mask = np .zeros (len (data ))
1158
- mask [deleted ] = 1
1159
- indexer = (ai - mask .cumsum ())[indexer ]
1160
-
1161
- offset = 0
1162
- for blk in agg_blocks :
1163
- loc = len (blk .mgr_locs )
1164
- blk .mgr_locs = indexer [offset : (offset + loc )]
1165
- offset += loc
1145
+ agg_items = data .reset_dropped_locs (agg_blocks , skipped )
1166
1146
1167
1147
return agg_blocks , agg_items
1168
1148
0 commit comments