Skip to content

Commit f638abd

Browse files
committed
Refactor to support more than one loc
1 parent e27d788 commit f638abd

File tree

1 file changed

+49
-6
lines changed

1 file changed

+49
-6
lines changed

pandas/core/internals/managers.py

+49-6
Original file line numberDiff line numberDiff line change
@@ -1235,19 +1235,62 @@ def value_getitem(placement):
12351235
+ self.refs[blkno_l + 1 :]
12361236
+ extra_refs
12371237
)
1238+
if len(blk_locs) > 1:
1239+
# Add the refs for the other deleted blocks
1240+
self.refs += [None] * (len(blk_locs) - 1)
1241+
1242+
# Fill before the first leftover block
1243+
nbs = []
1244+
leftover_start_i = leftover_blocks[0].mgr_locs.as_slice.start
1245+
1246+
if leftover_start_i != 0:
1247+
nbs.append(
1248+
new_block_2d(
1249+
np.tile(value, (leftover_start_i, 1)),
1250+
placement=BlockPlacement(slice(0, leftover_start_i)),
1251+
)
1252+
)
1253+
1254+
# Every hole in between leftover blocks is where we need to insert
1255+
# a new block
1256+
for i in range(len(leftover_blocks) - 1):
1257+
curr_block = leftover_blocks[i]
1258+
next_block = leftover_blocks[i + 1]
1259+
1260+
curr_end = curr_block.as_slice.end
1261+
next_start = next_block.as_slice.start
1262+
1263+
num_to_fill = next_start - (curr_end - 1)
1264+
nb = new_block_2d(
1265+
np.tile(value, (num_to_fill, 1)),
1266+
placement=BlockPlacement(
1267+
slice(curr_end, curr_end + num_to_fill)
1268+
),
1269+
)
1270+
nbs.append(nb)
1271+
1272+
# Fill after the last leftover block
1273+
last_del_loc = blk_locs[-1] + 1
1274+
last_leftover_loc = leftover_blocks[-1].mgr_locs.as_slice.stop
1275+
if last_del_loc > last_leftover_loc:
1276+
diff = last_del_loc - last_leftover_loc
1277+
nbs.append(
1278+
new_block_2d(
1279+
np.tile(value, (diff, 1)),
1280+
placement=BlockPlacement(
1281+
slice(last_leftover_loc, last_del_loc)
1282+
),
1283+
)
1284+
)
12381285

1239-
# TODO: Are we sure we want 2D?
1240-
# Also generalize this to bigger than 1-D locs
1241-
nb = new_block_2d(
1242-
value,
1243-
placement=BlockPlacement(slice(blk_locs[0], blk_locs[0] + 1)),
1244-
)
12451286
old_blocks = self.blocks
1287+
nb = nbs[0]
12461288
new_blocks = (
12471289
old_blocks[:blkno_l]
12481290
+ (nb,)
12491291
+ old_blocks[blkno_l + 1 :]
12501292
+ leftover_blocks
1293+
+ tuple(nbs)
12511294
)
12521295
self._blklocs[nb.mgr_locs.indexer] = np.arange(len(nb))
12531296
for i, nb in enumerate(leftover_blocks):

0 commit comments

Comments
 (0)