1
1
from __future__ import annotations
2
2
3
+ import copy
3
4
import itertools
4
5
from typing import (
5
6
TYPE_CHECKING ,
@@ -332,15 +333,12 @@ def _get_mgr_concatenation_plan(mgr: BlockManager, indexers: dict[int, np.ndarra
332
333
)
333
334
)
334
335
335
- if not unit_no_ax0_reindexing :
336
- # create block from subset of columns
337
- # Note: Blocks with only 1 column will always have unit_no_ax0_reindexing,
338
- # so we will never get here with ExtensionBlock.
339
- blk = blk . getitem_block ( ax0_blk_indexer )
336
+ # Omit indexer if no item reindexing is required.
337
+ if unit_no_ax0_reindexing :
338
+ join_unit_indexers . pop ( 0 , None )
339
+ else :
340
+ join_unit_indexers [ 0 ] = ax0_blk_indexer
340
341
341
- # Assertions disabled for performance
342
- # assert blk._mgr_locs.as_slice == placements.as_slice
343
- # assert blk.shape[0] == shape[0]
344
342
unit = JoinUnit (blk , shape , join_unit_indexers )
345
343
346
344
plan .append ((placements , unit ))
@@ -354,7 +352,6 @@ def __init__(self, block: Block, shape: Shape, indexers=None):
354
352
# Note: block is None implies indexers is None, but not vice-versa
355
353
if indexers is None :
356
354
indexers = {}
357
- # we should *never* have `0 in indexers`
358
355
self .block = block
359
356
self .indexers = indexers
360
357
self .shape = shape
@@ -606,11 +603,20 @@ def _trim_join_unit(join_unit: JoinUnit, length: int) -> JoinUnit:
606
603
607
604
Extra items that didn't fit are returned as a separate block.
608
605
"""
609
- assert 0 not in join_unit .indexers
610
- extra_indexers = join_unit .indexers
606
+ if 0 not in join_unit .indexers :
607
+ extra_indexers = join_unit .indexers
608
+
609
+ if join_unit .block is None :
610
+ extra_block = None
611
+ else :
612
+ extra_block = join_unit .block .getitem_block (slice (length , None ))
613
+ join_unit .block = join_unit .block .getitem_block (slice (length ))
614
+ else :
615
+ extra_block = join_unit .block
611
616
612
- extra_block = join_unit .block .getitem_block (slice (length , None ))
613
- join_unit .block = join_unit .block .getitem_block (slice (length ))
617
+ extra_indexers = copy .copy (join_unit .indexers )
618
+ extra_indexers [0 ] = extra_indexers [0 ][length :]
619
+ join_unit .indexers [0 ] = join_unit .indexers [0 ][:length ]
614
620
615
621
extra_shape = (join_unit .shape [0 ] - length ,) + join_unit .shape [1 :]
616
622
join_unit .shape = (length ,) + join_unit .shape [1 :]
0 commit comments