Skip to content

REF: de-duplicate take_nd #48339

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Sep 15, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 6 additions & 28 deletions pandas/core/internals/blocks.py
Original file line number Diff line number Diff line change
Expand Up @@ -881,8 +881,13 @@ def take_nd(
)

# Called from three places in managers, all of which satisfy
# this assertion
# these assertions
if isinstance(self, ExtensionBlock):
# NB: in this case, the 'axis' kwarg will be ignored in the
# algos.take_nd call above.
assert not (self.ndim == 1 and new_mgr_locs is None)
assert not (axis == 0 and new_mgr_locs is None)

if new_mgr_locs is None:
new_mgr_locs = self._mgr_locs

Expand Down Expand Up @@ -1752,33 +1757,6 @@ def is_view(self) -> bool:
def is_numeric(self):
return self.values.dtype._is_numeric

def take_nd(
self,
indexer: npt.NDArray[np.intp],
axis: int = 0,
new_mgr_locs: BlockPlacement | None = None,
fill_value=lib.no_default,
) -> Block:
"""
Take values according to indexer and return them as a block.
"""
if fill_value is lib.no_default:
fill_value = None

# TODO(EA2D): special case not needed with 2D EAs
# axis doesn't matter; we are really a single-dim object
# but are passed the axis depending on the calling routing
# if its REALLY axis 0, then this will be a reindex and not a take
new_values = self.values.take(indexer, fill_value=fill_value, allow_fill=True)

# Called from three places in managers, all of which satisfy
# this assertion
assert not (self.ndim == 1 and new_mgr_locs is None)
if new_mgr_locs is None:
new_mgr_locs = self._mgr_locs

return self.make_block_same_class(new_values, new_mgr_locs)

def _slice(
self, slicer: slice | npt.NDArray[np.bool_] | npt.NDArray[np.intp]
) -> ExtensionArray:
Expand Down
3 changes: 3 additions & 0 deletions pandas/core/internals/managers.py
Original file line number Diff line number Diff line change
Expand Up @@ -725,6 +725,9 @@ def reindex_indexer(
result.axes[axis] = new_axis
return result

# Should be intp, but in some cases we get int64 on 32bit builds
assert isinstance(indexer, np.ndarray)

# some axes don't allow reindexing with dups
if not allow_dups:
self.axes[axis]._validate_can_reindex(indexer)
Expand Down