Skip to content

Commit e4e0333

Browse files
jbrockmendelJulianWgs
authored andcommitted
PERF: use ndarray.take instead of algos.take (pandas-dev#40852)
1 parent 7563d06 commit e4e0333

File tree

5 files changed

+9
-17
lines changed

5 files changed

+9
-17
lines changed

pandas/core/groupby/ops.py

+2-3
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,6 @@
6767
maybe_fill,
6868
)
6969

70-
from pandas.core import algorithms
7170
from pandas.core.arrays import ExtensionArray
7271
import pandas.core.common as com
7372
from pandas.core.frame import DataFrame
@@ -757,7 +756,7 @@ def _aggregate_series_fast(self, obj: Series, func: F):
757756
# avoids object / Series creation overhead
758757
indexer = get_group_index_sorter(group_index, ngroups)
759758
obj = obj.take(indexer)
760-
group_index = algorithms.take_nd(group_index, indexer, allow_fill=False)
759+
group_index = group_index.take(indexer)
761760
grouper = libreduction.SeriesGrouper(obj, func, group_index, ngroups)
762761
result, counts = grouper.get_result()
763762
return result, counts
@@ -988,7 +987,7 @@ def __init__(self, data: FrameOrSeries, labels, ngroups: int, axis: int = 0):
988987
@cache_readonly
989988
def slabels(self) -> np.ndarray: # np.ndarray[np.intp]
990989
# Sorted labels
991-
return algorithms.take_nd(self.labels, self._sort_idx, allow_fill=False)
990+
return self.labels.take(self._sort_idx)
992991

993992
@cache_readonly
994993
def _sort_idx(self) -> np.ndarray: # np.ndarray[np.intp]

pandas/core/indexes/base.py

+2-4
Original file line numberDiff line numberDiff line change
@@ -2998,7 +2998,7 @@ def _union(self, other: Index, sort):
29982998
missing = algos.unique1d(self.get_indexer_non_unique(other)[1])
29992999

30003000
if len(missing) > 0:
3001-
other_diff = algos.take_nd(rvals, missing, allow_fill=False)
3001+
other_diff = rvals.take(missing)
30023002
result = concat_compat((lvals, other_diff))
30033003
else:
30043004
# error: Incompatible types in assignment (expression has type
@@ -4237,9 +4237,7 @@ def _get_leaf_sorter(labels: list[np.ndarray]) -> np.ndarray:
42374237
)
42384238

42394239
if right_lev_indexer is not None:
4240-
right_indexer = algos.take_nd(
4241-
right_lev_indexer, join_index.codes[level], allow_fill=False
4242-
)
4240+
right_indexer = right_lev_indexer.take(join_index.codes[level])
42434241
else:
42444242
right_indexer = join_index.codes[level]
42454243

pandas/core/indexes/multi.py

+2-6
Original file line numberDiff line numberDiff line change
@@ -3533,14 +3533,10 @@ def equals(self, other: object) -> bool:
35333533
if not np.array_equal(self_mask, other_mask):
35343534
return False
35353535
self_codes = self_codes[~self_mask]
3536-
self_values = algos.take_nd(
3537-
np.asarray(self.levels[i]._values), self_codes, allow_fill=False
3538-
)
3536+
self_values = self.levels[i]._values.take(self_codes)
35393537

35403538
other_codes = other_codes[~other_mask]
3541-
other_values = other_values = algos.take_nd(
3542-
np.asarray(other.levels[i]._values), other_codes, allow_fill=False
3543-
)
3539+
other_values = other.levels[i]._values.take(other_codes)
35443540

35453541
# since we use NaT both datetime64 and timedelta64 we can have a
35463542
# situation where a level is typed say timedelta64 in self (IOW it

pandas/core/internals/managers.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -263,7 +263,7 @@ def items(self) -> Index:
263263

264264
def get_dtypes(self):
265265
dtypes = np.array([blk.dtype for blk in self.blocks])
266-
return algos.take_nd(dtypes, self.blknos, allow_fill=False)
266+
return dtypes.take(self.blknos)
267267

268268
@property
269269
def arrays(self) -> list[ArrayLike]:

pandas/core/sorting.py

+2-3
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@
3131
)
3232
from pandas.core.dtypes.missing import isna
3333

34-
from pandas.core import algorithms
3534
from pandas.core.construction import extract_array
3635

3736
if TYPE_CHECKING:
@@ -668,10 +667,10 @@ def _reorder_by_uniques(uniques, labels):
668667
mask = labels < 0
669668

670669
# move labels to right locations (ie, unsort ascending labels)
671-
labels = algorithms.take_nd(reverse_indexer, labels, allow_fill=False)
670+
labels = reverse_indexer.take(labels)
672671
np.putmask(labels, mask, -1)
673672

674673
# sort observed ids
675-
uniques = algorithms.take_nd(uniques, sorter, allow_fill=False)
674+
uniques = uniques.take(sorter)
676675

677676
return uniques, labels

0 commit comments

Comments
 (0)