Skip to content

Commit faa8114

Browse files
authored
REF: de-duplicate take_nd (#48339)
* REF: de-duplicate take_nd * revert * fix 32bit
1 parent 752fda8 commit faa8114

File tree

2 files changed

+9
-28
lines changed

2 files changed

+9
-28
lines changed

pandas/core/internals/blocks.py

+6-28
Original file line numberDiff line numberDiff line change
@@ -881,8 +881,13 @@ def take_nd(
881881
)
882882

883883
# Called from three places in managers, all of which satisfy
884-
# this assertion
884+
# these assertions
885+
if isinstance(self, ExtensionBlock):
886+
# NB: in this case, the 'axis' kwarg will be ignored in the
887+
# algos.take_nd call above.
888+
assert not (self.ndim == 1 and new_mgr_locs is None)
885889
assert not (axis == 0 and new_mgr_locs is None)
890+
886891
if new_mgr_locs is None:
887892
new_mgr_locs = self._mgr_locs
888893

@@ -1753,33 +1758,6 @@ def is_view(self) -> bool:
17531758
def is_numeric(self):
17541759
return self.values.dtype._is_numeric
17551760

1756-
def take_nd(
1757-
self,
1758-
indexer: npt.NDArray[np.intp],
1759-
axis: int = 0,
1760-
new_mgr_locs: BlockPlacement | None = None,
1761-
fill_value=lib.no_default,
1762-
) -> Block:
1763-
"""
1764-
Take values according to indexer and return them as a block.
1765-
"""
1766-
if fill_value is lib.no_default:
1767-
fill_value = None
1768-
1769-
# TODO(EA2D): special case not needed with 2D EAs
1770-
# axis doesn't matter; we are really a single-dim object
1771-
# but are passed the axis depending on the calling routing
1772-
# if its REALLY axis 0, then this will be a reindex and not a take
1773-
new_values = self.values.take(indexer, fill_value=fill_value, allow_fill=True)
1774-
1775-
# Called from three places in managers, all of which satisfy
1776-
# this assertion
1777-
assert not (self.ndim == 1 and new_mgr_locs is None)
1778-
if new_mgr_locs is None:
1779-
new_mgr_locs = self._mgr_locs
1780-
1781-
return self.make_block_same_class(new_values, new_mgr_locs)
1782-
17831761
def _slice(
17841762
self, slicer: slice | npt.NDArray[np.bool_] | npt.NDArray[np.intp]
17851763
) -> ExtensionArray:

pandas/core/internals/managers.py

+3
Original file line numberDiff line numberDiff line change
@@ -726,6 +726,9 @@ def reindex_indexer(
726726
result.axes[axis] = new_axis
727727
return result
728728

729+
# Should be intp, but in some cases we get int64 on 32bit builds
730+
assert isinstance(indexer, np.ndarray)
731+
729732
# some axes don't allow reindexing with dups
730733
if not allow_dups:
731734
self.axes[axis]._validate_can_reindex(indexer)

0 commit comments

Comments
 (0)