Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3import functools
4from typing import (
5 TYPE_CHECKING,
6 cast,
7 overload,
8)
10import numpy as np
12from pandas._libs import (
13 algos as libalgos,
14 lib,
15)
17from pandas.core.dtypes.cast import maybe_promote
18from pandas.core.dtypes.common import (
19 ensure_platform_int,
20 is_1d_only_ea_dtype,
21)
22from pandas.core.dtypes.missing import na_value_for_dtype
24from pandas.core.construction import ensure_wrapped_if_datetimelike
26if TYPE_CHECKING:
27 from pandas._typing import (
28 ArrayLike,
29 AxisInt,
30 npt,
31 )
33 from pandas.core.arrays._mixins import NDArrayBackedExtensionArray
34 from pandas.core.arrays.base import ExtensionArray
37@overload
38def take_nd(
39 arr: np.ndarray,
40 indexer,
41 axis: AxisInt = ...,
42 fill_value=...,
43 allow_fill: bool = ...,
44) -> np.ndarray: ...
47@overload
48def take_nd(
49 arr: ExtensionArray,
50 indexer,
51 axis: AxisInt = ...,
52 fill_value=...,
53 allow_fill: bool = ...,
54) -> ArrayLike: ...
57def take_nd(
58 arr: ArrayLike,
59 indexer,
60 axis: AxisInt = 0,
61 fill_value=lib.no_default,
62 allow_fill: bool = True,
63) -> ArrayLike:
64 """
65 Specialized Cython take which sets NaN values in one pass
67 This dispatches to ``take`` defined on ExtensionArrays.
69 Note: this function assumes that the indexer is a valid(ated) indexer with
70 no out of bound indices.
72 Parameters
73 ----------
74 arr : np.ndarray or ExtensionArray
75 Input array.
76 indexer : ndarray
77 1-D array of indices to take, subarrays corresponding to -1 value
78 indices are filed with fill_value
79 axis : int, default 0
80 Axis to take from
81 fill_value : any, default np.nan
82 Fill value to replace -1 values with
83 allow_fill : bool, default True
84 If False, indexer is assumed to contain no -1 values so no filling
85 will be done. This short-circuits computation of a mask. Result is
86 undefined if allow_fill == False and -1 is present in indexer.
88 Returns
89 -------
90 subarray : np.ndarray or ExtensionArray
91 May be the same type as the input, or cast to an ndarray.
92 """
93 if fill_value is lib.no_default:
94 fill_value = na_value_for_dtype(arr.dtype, compat=False)
95 elif lib.is_np_dtype(arr.dtype, "mM"):
96 dtype, fill_value = maybe_promote(arr.dtype, fill_value)
97 if arr.dtype != dtype:
98 # EA.take is strict about returning a new object of the same type
99 # so for that case cast upfront
100 arr = arr.astype(dtype)
102 if not isinstance(arr, np.ndarray):
103 # i.e. ExtensionArray,
104 # includes for EA to catch DatetimeArray, TimedeltaArray
105 if not is_1d_only_ea_dtype(arr.dtype):
106 # i.e. DatetimeArray, TimedeltaArray
107 arr = cast("NDArrayBackedExtensionArray", arr)
108 return arr.take(
109 indexer, fill_value=fill_value, allow_fill=allow_fill, axis=axis
110 )
112 return arr.take(indexer, fill_value=fill_value, allow_fill=allow_fill)
114 arr = np.asarray(arr)
115 return _take_nd_ndarray(arr, indexer, axis, fill_value, allow_fill)
118def _take_nd_ndarray(
119 arr: np.ndarray,
120 indexer: npt.NDArray[np.intp] | None,
121 axis: AxisInt,
122 fill_value,
123 allow_fill: bool,
124) -> np.ndarray:
125 if indexer is None:
126 indexer = np.arange(arr.shape[axis], dtype=np.intp)
127 dtype, fill_value = arr.dtype, arr.dtype.type()
128 else:
129 indexer = ensure_platform_int(indexer)
131 dtype, fill_value, mask_info = _take_preprocess_indexer_and_fill_value(
132 arr, indexer, fill_value, allow_fill
133 )
135 flip_order = False
136 if arr.ndim == 2 and arr.flags.f_contiguous:
137 flip_order = True
139 if flip_order:
140 arr = arr.T
141 axis = arr.ndim - axis - 1
143 # at this point, it's guaranteed that dtype can hold both the arr values
144 # and the fill_value
145 out_shape_ = list(arr.shape)
146 out_shape_[axis] = len(indexer)
147 out_shape = tuple(out_shape_)
148 if arr.flags.f_contiguous and axis == arr.ndim - 1:
149 # minor tweak that can make an order-of-magnitude difference
150 # for dataframes initialized directly from 2-d ndarrays
151 # (s.t. df.values is c-contiguous and df._mgr.blocks[0] is its
152 # f-contiguous transpose)
153 out = np.empty(out_shape, dtype=dtype, order="F")
154 else:
155 out = np.empty(out_shape, dtype=dtype)
157 func = _get_take_nd_function(
158 arr.ndim, arr.dtype, out.dtype, axis=axis, mask_info=mask_info
159 )
160 func(arr, indexer, out, fill_value)
162 if flip_order:
163 out = out.T
164 return out
167def take_2d_multi(
168 arr: np.ndarray,
169 indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]],
170 fill_value=np.nan,
171) -> np.ndarray:
172 """
173 Specialized Cython take which sets NaN values in one pass.
174 """
175 # This is only called from one place in DataFrame._reindex_multi,
176 # so we know indexer is well-behaved.
177 assert indexer is not None
178 assert indexer[0] is not None
179 assert indexer[1] is not None
181 row_idx, col_idx = indexer
183 row_idx = ensure_platform_int(row_idx)
184 col_idx = ensure_platform_int(col_idx)
185 indexer = row_idx, col_idx
186 mask_info = None
188 # check for promotion based on types only (do this first because
189 # it's faster than computing a mask)
190 dtype, fill_value = maybe_promote(arr.dtype, fill_value)
191 if dtype != arr.dtype:
192 # check if promotion is actually required based on indexer
193 row_mask = row_idx == -1
194 col_mask = col_idx == -1
195 row_needs = row_mask.any()
196 col_needs = col_mask.any()
197 mask_info = (row_mask, col_mask), (row_needs, col_needs)
199 if not (row_needs or col_needs):
200 # if not, then depromote, set fill_value to dummy
201 # (it won't be used but we don't want the cython code
202 # to crash when trying to cast it to dtype)
203 dtype, fill_value = arr.dtype, arr.dtype.type()
205 # at this point, it's guaranteed that dtype can hold both the arr values
206 # and the fill_value
207 out_shape = len(row_idx), len(col_idx)
208 out = np.empty(out_shape, dtype=dtype)
210 func = _take_2d_multi_dict.get((arr.dtype.name, out.dtype.name), None)
211 if func is None and arr.dtype != out.dtype:
212 func = _take_2d_multi_dict.get((out.dtype.name, out.dtype.name), None)
213 if func is not None:
214 func = _convert_wrapper(func, out.dtype)
216 if func is not None:
217 func(arr, indexer, out=out, fill_value=fill_value)
218 else:
219 # test_reindex_multi
220 _take_2d_multi_object(
221 arr, indexer, out, fill_value=fill_value, mask_info=mask_info
222 )
224 return out
227@functools.lru_cache
228def _get_take_nd_function_cached(
229 ndim: int, arr_dtype: np.dtype, out_dtype: np.dtype, axis: AxisInt
230):
231 """
232 Part of _get_take_nd_function below that doesn't need `mask_info` and thus
233 can be cached (mask_info potentially contains a numpy ndarray which is not
234 hashable and thus cannot be used as argument for cached function).
235 """
236 tup = (arr_dtype.name, out_dtype.name)
237 if ndim == 1:
238 func = _take_1d_dict.get(tup, None)
239 elif ndim == 2:
240 if axis == 0:
241 func = _take_2d_axis0_dict.get(tup, None)
242 else:
243 func = _take_2d_axis1_dict.get(tup, None)
244 if func is not None:
245 return func
247 # We get here with string, uint, float16, and complex dtypes that could
248 # potentially be handled in algos_take_helper.
249 # Also a couple with (M8[ns], object) and (m8[ns], object)
250 tup = (out_dtype.name, out_dtype.name)
251 if ndim == 1:
252 func = _take_1d_dict.get(tup, None)
253 elif ndim == 2:
254 if axis == 0:
255 func = _take_2d_axis0_dict.get(tup, None)
256 else:
257 func = _take_2d_axis1_dict.get(tup, None)
258 if func is not None:
259 func = _convert_wrapper(func, out_dtype)
260 return func
262 return None
265def _get_take_nd_function(
266 ndim: int,
267 arr_dtype: np.dtype,
268 out_dtype: np.dtype,
269 axis: AxisInt = 0,
270 mask_info=None,
271):
272 """
273 Get the appropriate "take" implementation for the given dimension, axis
274 and dtypes.
275 """
276 func = None
277 if ndim <= 2:
278 # for this part we don't need `mask_info` -> use the cached algo lookup
279 func = _get_take_nd_function_cached(ndim, arr_dtype, out_dtype, axis)
281 if func is None:
283 def func(arr, indexer, out, fill_value=np.nan) -> None:
284 indexer = ensure_platform_int(indexer)
285 _take_nd_object(
286 arr, indexer, out, axis=axis, fill_value=fill_value, mask_info=mask_info
287 )
289 return func
292def _view_wrapper(f, arr_dtype=None, out_dtype=None, fill_wrap=None):
293 def wrapper(
294 arr: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=np.nan
295 ) -> None:
296 if arr_dtype is not None:
297 arr = arr.view(arr_dtype)
298 if out_dtype is not None:
299 out = out.view(out_dtype)
300 if fill_wrap is not None:
301 # FIXME: if we get here with dt64/td64 we need to be sure we have
302 # matching resos
303 if fill_value.dtype.kind == "m":
304 fill_value = fill_value.astype("m8[ns]")
305 else:
306 fill_value = fill_value.astype("M8[ns]")
307 fill_value = fill_wrap(fill_value)
309 f(arr, indexer, out, fill_value=fill_value)
311 return wrapper
314def _convert_wrapper(f, conv_dtype):
315 def wrapper(
316 arr: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=np.nan
317 ) -> None:
318 if conv_dtype == object:
319 # GH#39755 avoid casting dt64/td64 to integers
320 arr = ensure_wrapped_if_datetimelike(arr)
321 arr = arr.astype(conv_dtype)
322 f(arr, indexer, out, fill_value=fill_value)
324 return wrapper
327_take_1d_dict = {
328 ("int8", "int8"): libalgos.take_1d_int8_int8,
329 ("int8", "int32"): libalgos.take_1d_int8_int32,
330 ("int8", "int64"): libalgos.take_1d_int8_int64,
331 ("int8", "float64"): libalgos.take_1d_int8_float64,
332 ("int16", "int16"): libalgos.take_1d_int16_int16,
333 ("int16", "int32"): libalgos.take_1d_int16_int32,
334 ("int16", "int64"): libalgos.take_1d_int16_int64,
335 ("int16", "float64"): libalgos.take_1d_int16_float64,
336 ("int32", "int32"): libalgos.take_1d_int32_int32,
337 ("int32", "int64"): libalgos.take_1d_int32_int64,
338 ("int32", "float64"): libalgos.take_1d_int32_float64,
339 ("int64", "int64"): libalgos.take_1d_int64_int64,
340 ("uint8", "uint8"): libalgos.take_1d_bool_bool,
341 ("uint16", "int64"): libalgos.take_1d_uint16_uint16,
342 ("uint32", "int64"): libalgos.take_1d_uint32_uint32,
343 ("uint64", "int64"): libalgos.take_1d_uint64_uint64,
344 ("int64", "float64"): libalgos.take_1d_int64_float64,
345 ("float32", "float32"): libalgos.take_1d_float32_float32,
346 ("float32", "float64"): libalgos.take_1d_float32_float64,
347 ("float64", "float64"): libalgos.take_1d_float64_float64,
348 ("object", "object"): libalgos.take_1d_object_object,
349 ("bool", "bool"): _view_wrapper(libalgos.take_1d_bool_bool, np.uint8, np.uint8),
350 ("bool", "object"): _view_wrapper(libalgos.take_1d_bool_object, np.uint8, None),
351 ("datetime64[ns]", "datetime64[ns]"): _view_wrapper(
352 libalgos.take_1d_int64_int64, np.int64, np.int64, np.int64
353 ),
354 ("timedelta64[ns]", "timedelta64[ns]"): _view_wrapper(
355 libalgos.take_1d_int64_int64, np.int64, np.int64, np.int64
356 ),
357}
359_take_2d_axis0_dict = {
360 ("int8", "int8"): libalgos.take_2d_axis0_int8_int8,
361 ("int8", "int32"): libalgos.take_2d_axis0_int8_int32,
362 ("int8", "int64"): libalgos.take_2d_axis0_int8_int64,
363 ("int8", "float64"): libalgos.take_2d_axis0_int8_float64,
364 ("int16", "int16"): libalgos.take_2d_axis0_int16_int16,
365 ("int16", "int32"): libalgos.take_2d_axis0_int16_int32,
366 ("int16", "int64"): libalgos.take_2d_axis0_int16_int64,
367 ("int16", "float64"): libalgos.take_2d_axis0_int16_float64,
368 ("int32", "int32"): libalgos.take_2d_axis0_int32_int32,
369 ("int32", "int64"): libalgos.take_2d_axis0_int32_int64,
370 ("int32", "float64"): libalgos.take_2d_axis0_int32_float64,
371 ("int64", "int64"): libalgos.take_2d_axis0_int64_int64,
372 ("int64", "float64"): libalgos.take_2d_axis0_int64_float64,
373 ("uint8", "uint8"): libalgos.take_2d_axis0_bool_bool,
374 ("uint16", "uint16"): libalgos.take_2d_axis0_uint16_uint16,
375 ("uint32", "uint32"): libalgos.take_2d_axis0_uint32_uint32,
376 ("uint64", "uint64"): libalgos.take_2d_axis0_uint64_uint64,
377 ("float32", "float32"): libalgos.take_2d_axis0_float32_float32,
378 ("float32", "float64"): libalgos.take_2d_axis0_float32_float64,
379 ("float64", "float64"): libalgos.take_2d_axis0_float64_float64,
380 ("object", "object"): libalgos.take_2d_axis0_object_object,
381 ("bool", "bool"): _view_wrapper(
382 libalgos.take_2d_axis0_bool_bool, np.uint8, np.uint8
383 ),
384 ("bool", "object"): _view_wrapper(
385 libalgos.take_2d_axis0_bool_object, np.uint8, None
386 ),
387 ("datetime64[ns]", "datetime64[ns]"): _view_wrapper(
388 libalgos.take_2d_axis0_int64_int64, np.int64, np.int64, fill_wrap=np.int64
389 ),
390 ("timedelta64[ns]", "timedelta64[ns]"): _view_wrapper(
391 libalgos.take_2d_axis0_int64_int64, np.int64, np.int64, fill_wrap=np.int64
392 ),
393}
395_take_2d_axis1_dict = {
396 ("int8", "int8"): libalgos.take_2d_axis1_int8_int8,
397 ("int8", "int32"): libalgos.take_2d_axis1_int8_int32,
398 ("int8", "int64"): libalgos.take_2d_axis1_int8_int64,
399 ("int8", "float64"): libalgos.take_2d_axis1_int8_float64,
400 ("int16", "int16"): libalgos.take_2d_axis1_int16_int16,
401 ("int16", "int32"): libalgos.take_2d_axis1_int16_int32,
402 ("int16", "int64"): libalgos.take_2d_axis1_int16_int64,
403 ("int16", "float64"): libalgos.take_2d_axis1_int16_float64,
404 ("int32", "int32"): libalgos.take_2d_axis1_int32_int32,
405 ("int32", "int64"): libalgos.take_2d_axis1_int32_int64,
406 ("int32", "float64"): libalgos.take_2d_axis1_int32_float64,
407 ("int64", "int64"): libalgos.take_2d_axis1_int64_int64,
408 ("int64", "float64"): libalgos.take_2d_axis1_int64_float64,
409 ("uint8", "uint8"): libalgos.take_2d_axis1_bool_bool,
410 ("uint16", "uint16"): libalgos.take_2d_axis1_uint16_uint16,
411 ("uint32", "uint32"): libalgos.take_2d_axis1_uint32_uint32,
412 ("uint64", "uint64"): libalgos.take_2d_axis1_uint64_uint64,
413 ("float32", "float32"): libalgos.take_2d_axis1_float32_float32,
414 ("float32", "float64"): libalgos.take_2d_axis1_float32_float64,
415 ("float64", "float64"): libalgos.take_2d_axis1_float64_float64,
416 ("object", "object"): libalgos.take_2d_axis1_object_object,
417 ("bool", "bool"): _view_wrapper(
418 libalgos.take_2d_axis1_bool_bool, np.uint8, np.uint8
419 ),
420 ("bool", "object"): _view_wrapper(
421 libalgos.take_2d_axis1_bool_object, np.uint8, None
422 ),
423 ("datetime64[ns]", "datetime64[ns]"): _view_wrapper(
424 libalgos.take_2d_axis1_int64_int64, np.int64, np.int64, fill_wrap=np.int64
425 ),
426 ("timedelta64[ns]", "timedelta64[ns]"): _view_wrapper(
427 libalgos.take_2d_axis1_int64_int64, np.int64, np.int64, fill_wrap=np.int64
428 ),
429}
431_take_2d_multi_dict = {
432 ("int8", "int8"): libalgos.take_2d_multi_int8_int8,
433 ("int8", "int32"): libalgos.take_2d_multi_int8_int32,
434 ("int8", "int64"): libalgos.take_2d_multi_int8_int64,
435 ("int8", "float64"): libalgos.take_2d_multi_int8_float64,
436 ("int16", "int16"): libalgos.take_2d_multi_int16_int16,
437 ("int16", "int32"): libalgos.take_2d_multi_int16_int32,
438 ("int16", "int64"): libalgos.take_2d_multi_int16_int64,
439 ("int16", "float64"): libalgos.take_2d_multi_int16_float64,
440 ("int32", "int32"): libalgos.take_2d_multi_int32_int32,
441 ("int32", "int64"): libalgos.take_2d_multi_int32_int64,
442 ("int32", "float64"): libalgos.take_2d_multi_int32_float64,
443 ("int64", "int64"): libalgos.take_2d_multi_int64_int64,
444 ("int64", "float64"): libalgos.take_2d_multi_int64_float64,
445 ("float32", "float32"): libalgos.take_2d_multi_float32_float32,
446 ("float32", "float64"): libalgos.take_2d_multi_float32_float64,
447 ("float64", "float64"): libalgos.take_2d_multi_float64_float64,
448 ("object", "object"): libalgos.take_2d_multi_object_object,
449 ("bool", "bool"): _view_wrapper(
450 libalgos.take_2d_multi_bool_bool, np.uint8, np.uint8
451 ),
452 ("bool", "object"): _view_wrapper(
453 libalgos.take_2d_multi_bool_object, np.uint8, None
454 ),
455 ("datetime64[ns]", "datetime64[ns]"): _view_wrapper(
456 libalgos.take_2d_multi_int64_int64, np.int64, np.int64, fill_wrap=np.int64
457 ),
458 ("timedelta64[ns]", "timedelta64[ns]"): _view_wrapper(
459 libalgos.take_2d_multi_int64_int64, np.int64, np.int64, fill_wrap=np.int64
460 ),
461}
464def _take_nd_object(
465 arr: np.ndarray,
466 indexer: npt.NDArray[np.intp],
467 out: np.ndarray,
468 axis: AxisInt,
469 fill_value,
470 mask_info,
471) -> None:
472 if mask_info is not None:
473 mask, needs_masking = mask_info
474 else:
475 mask = indexer == -1
476 needs_masking = mask.any()
477 if arr.dtype != out.dtype:
478 arr = arr.astype(out.dtype)
479 if arr.shape[axis] > 0:
480 arr.take(indexer, axis=axis, out=out)
481 if needs_masking:
482 outindexer = [slice(None)] * arr.ndim
483 outindexer[axis] = mask
484 out[tuple(outindexer)] = fill_value
487def _take_2d_multi_object(
488 arr: np.ndarray,
489 indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]],
490 out: np.ndarray,
491 fill_value,
492 mask_info,
493) -> None:
494 # this is not ideal, performance-wise, but it's better than raising
495 # an exception (best to optimize in Cython to avoid getting here)
496 row_idx, col_idx = indexer # both np.intp
497 if mask_info is not None:
498 (row_mask, col_mask), (row_needs, col_needs) = mask_info
499 else:
500 row_mask = row_idx == -1
501 col_mask = col_idx == -1
502 row_needs = row_mask.any()
503 col_needs = col_mask.any()
504 if fill_value is not None:
505 if row_needs:
506 out[row_mask, :] = fill_value
507 if col_needs:
508 out[:, col_mask] = fill_value
509 for i, u_ in enumerate(row_idx):
510 if u_ != -1:
511 for j, v in enumerate(col_idx):
512 if v != -1:
513 out[i, j] = arr[u_, v]
516def _take_preprocess_indexer_and_fill_value(
517 arr: np.ndarray,
518 indexer: npt.NDArray[np.intp],
519 fill_value,
520 allow_fill: bool,
521 mask: npt.NDArray[np.bool_] | None = None,
522):
523 mask_info: tuple[np.ndarray | None, bool] | None = None
525 if not allow_fill:
526 dtype, fill_value = arr.dtype, arr.dtype.type()
527 mask_info = None, False
528 else:
529 # check for promotion based on types only (do this first because
530 # it's faster than computing a mask)
531 dtype, fill_value = maybe_promote(arr.dtype, fill_value)
532 if dtype != arr.dtype:
533 # check if promotion is actually required based on indexer
534 if mask is not None:
535 needs_masking = True
536 else:
537 mask = indexer == -1
538 needs_masking = bool(mask.any())
539 mask_info = mask, needs_masking
540 if not needs_masking:
541 # if not, then depromote, set fill_value to dummy
542 # (it won't be used but we don't want the cython code
543 # to crash when trying to cast it to dtype)
544 dtype, fill_value = arr.dtype, arr.dtype.type()
546 return dtype, fill_value, mask_info