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
3from functools import wraps
4from typing import (
5 TYPE_CHECKING,
6 Any,
7 Literal,
8 Self,
9 cast,
10 overload,
11)
13import numpy as np
15from pandas._libs import lib
16from pandas._libs.arrays import NDArrayBacked
17from pandas._libs.tslibs import is_supported_dtype
18from pandas.errors import AbstractMethodError
19from pandas.util._validators import (
20 validate_bool_kwarg,
21 validate_insert_loc,
22)
24from pandas.core.dtypes.cast import construct_1d_object_array_from_listlike
25from pandas.core.dtypes.common import pandas_dtype
26from pandas.core.dtypes.dtypes import (
27 DatetimeTZDtype,
28 ExtensionDtype,
29 PeriodDtype,
30)
31from pandas.core.dtypes.missing import array_equivalent
33from pandas.core import missing
34from pandas.core.algorithms import (
35 take,
36 unique,
37 value_counts_internal as value_counts,
38)
39from pandas.core.array_algos.quantile import quantile_with_mask
40from pandas.core.array_algos.transforms import shift
41from pandas.core.arrays.base import ExtensionArray
42from pandas.core.construction import extract_array
43from pandas.core.indexers import (
44 check_array_indexer,
45 getitem_returns_view,
46)
47from pandas.core.sorting import nargminmax
49if TYPE_CHECKING:
50 from collections.abc import Sequence
52 from pandas._typing import (
53 ArrayLike,
54 AxisInt,
55 Dtype,
56 F,
57 FillnaOptions,
58 NumpySorter,
59 NumpyValueArrayLike,
60 PositionalIndexer2D,
61 PositionalIndexerTuple,
62 ScalarIndexer,
63 SequenceIndexer,
64 Shape,
65 SortKind,
66 TakeIndexer,
67 npt,
68 )
70 from pandas import Series
73def ravel_compat(meth: F) -> F:
74 """
75 Decorator to ravel a 2D array before passing it to a cython operation,
76 then reshape the result to our own shape.
77 """
79 @wraps(meth)
80 def method(self, *args, **kwargs):
81 if self.ndim == 1:
82 return meth(self, *args, **kwargs)
84 flags = self._ndarray.flags
85 flat = self.ravel("K")
86 result = meth(flat, *args, **kwargs)
87 order = "F" if flags.f_contiguous else "C"
88 return result.reshape(self.shape, order=order)
90 return cast("F", method)
93class NDArrayBackedExtensionArray(NDArrayBacked, ExtensionArray):
94 """
95 ExtensionArray that is backed by a single NumPy ndarray.
96 """
98 _ndarray: np.ndarray
100 # scalar used to denote NA value inside our self._ndarray, e.g. -1
101 # for Categorical, iNaT for Period. Outside of object dtype,
102 # self.isna() should be exactly locations in self._ndarray with
103 # _internal_fill_value.
104 _internal_fill_value: Any
106 def _box_func(self, x):
107 """
108 Wrap numpy type in our dtype.type if necessary.
109 """
110 return x
112 def _validate_scalar(self, value):
113 # used by NDArrayBackedExtensionIndex.insert
114 raise AbstractMethodError(self)
116 # ------------------------------------------------------------------------
118 @overload
119 def view(self, dtype: None = ...) -> Self: ...
121 @overload
122 def view(self, dtype: Dtype | None = ...) -> ArrayLike: ...
124 def view(self, dtype: Dtype | None = None) -> ArrayLike:
125 # We handle datetime64, datetime64tz, timedelta64, and period
126 # dtypes here. Everything else we pass through to the underlying
127 # ndarray.
128 if dtype is None or dtype is self.dtype:
129 return self._from_backing_data(self._ndarray)
131 if isinstance(dtype, type):
132 # we sometimes pass non-dtype objects, e.g np.ndarray;
133 # pass those through to the underlying ndarray
134 return self._ndarray.view(dtype)
136 dtype = pandas_dtype(dtype)
137 arr = self._ndarray
139 if isinstance(dtype, PeriodDtype):
140 cls = dtype.construct_array_type()
141 return cls(arr.view("i8"), dtype=dtype)
142 elif isinstance(dtype, DatetimeTZDtype):
143 dt_cls = dtype.construct_array_type()
144 dt64_values = arr.view(f"M8[{dtype.unit}]")
145 return dt_cls._simple_new(dt64_values, dtype=dtype)
146 elif lib.is_np_dtype(dtype, "M") and is_supported_dtype(dtype):
147 from pandas.core.arrays import DatetimeArray
149 dt64_values = arr.view(dtype)
150 return DatetimeArray._simple_new(dt64_values, dtype=dtype)
151 elif lib.is_np_dtype(dtype, "m") and is_supported_dtype(dtype):
152 from pandas.core.arrays import TimedeltaArray
154 td64_values = arr.view(dtype)
155 return TimedeltaArray._simple_new(td64_values, dtype=dtype)
156 # error: Argument "dtype" to "view" of "ndarray" has incompatible type
157 # "ExtensionDtype | dtype[Any]"; expected "dtype[Any] | _HasDType[dtype[Any]]"
158 return arr.view(dtype=dtype) # type: ignore[arg-type]
160 def take(
161 self,
162 indices: TakeIndexer,
163 *,
164 allow_fill: bool = False,
165 fill_value: Any = None,
166 axis: AxisInt = 0,
167 ) -> Self:
168 if allow_fill:
169 fill_value = self._validate_scalar(fill_value)
171 new_data = take(
172 self._ndarray,
173 indices,
174 allow_fill=allow_fill,
175 fill_value=fill_value,
176 axis=axis,
177 )
178 return self._from_backing_data(new_data)
180 # ------------------------------------------------------------------------
182 def equals(self, other) -> bool:
183 if type(self) is not type(other):
184 return False
185 if self.dtype != other.dtype:
186 return False
187 return bool(array_equivalent(self._ndarray, other._ndarray, dtype_equal=True))
189 @classmethod
190 def _from_factorized(cls, values, original):
191 assert values.dtype == original._ndarray.dtype
192 return original._from_backing_data(values)
194 def _values_for_argsort(self) -> np.ndarray:
195 return self._ndarray
197 def _values_for_factorize(self):
198 return self._ndarray, self._internal_fill_value
200 def _hash_pandas_object(
201 self, *, encoding: str, hash_key: str, categorize: bool
202 ) -> npt.NDArray[np.uint64]:
203 from pandas.core.util.hashing import hash_array
205 values = self._ndarray
206 return hash_array(
207 values, encoding=encoding, hash_key=hash_key, categorize=categorize
208 )
210 def _cast_pointwise_result(self, values: ArrayLike) -> ArrayLike:
211 if not (isinstance(values, np.ndarray) and values.dtype == object):
212 values = construct_1d_object_array_from_listlike(values) # type: ignore[arg-type]
213 return lib.maybe_convert_objects(values, convert_non_numeric=True)
215 # Signature of "argmin" incompatible with supertype "ExtensionArray"
216 def argmin(self, axis: AxisInt = 0, skipna: bool = True): # type: ignore[override]
217 # override base class by adding axis keyword
218 validate_bool_kwarg(skipna, "skipna")
219 if not skipna and self._hasna:
220 raise ValueError("Encountered an NA value with skipna=False")
221 return nargminmax(self, "argmin", axis=axis)
223 # Signature of "argmax" incompatible with supertype "ExtensionArray"
224 def argmax(self, axis: AxisInt = 0, skipna: bool = True): # type: ignore[override]
225 # override base class by adding axis keyword
226 validate_bool_kwarg(skipna, "skipna")
227 if not skipna and self._hasna:
228 raise ValueError("Encountered an NA value with skipna=False")
229 return nargminmax(self, "argmax", axis=axis)
231 def unique(self) -> Self:
232 new_data = unique(self._ndarray)
233 return self._from_backing_data(new_data)
235 def sort(
236 self,
237 *,
238 ascending: bool = True,
239 kind: SortKind = "quicksort",
240 na_position: str = "last",
241 ) -> None:
242 if self._readonly:
243 raise ValueError("Cannot modify read-only array")
244 sort_indices = self.argsort(
245 ascending=ascending, kind=kind, na_position=na_position
246 )
247 self._ndarray[:] = self._ndarray[sort_indices]
249 @classmethod
250 def _concat_same_type(
251 cls,
252 to_concat: Sequence[Self],
253 axis: AxisInt = 0,
254 ) -> Self:
255 """
256 Concatenate multiple arrays of this dtype.
258 Parameters
259 ----------
260 to_concat : sequence of this type
262 Returns
263 -------
264 ExtensionArray
265 """
266 if not lib.dtypes_all_equal([x.dtype for x in to_concat]):
267 dtypes = {str(x.dtype) for x in to_concat}
268 raise ValueError("to_concat must have the same dtype", dtypes)
270 return super()._concat_same_type(to_concat, axis=axis)
272 def searchsorted(
273 self,
274 value: NumpyValueArrayLike | ExtensionArray,
275 side: Literal["left", "right"] = "left",
276 sorter: NumpySorter | None = None,
277 ) -> npt.NDArray[np.intp] | np.intp:
278 """
279 Find indices where elements should be inserted to maintain order.
281 Find the indices into a sorted array `self` (a) such that, if the
282 corresponding elements in `value` were inserted before the indices,
283 the order of `self` would be preserved.
285 Assuming that `self` is sorted:
287 ====== ================================
288 `side` returned index `i` satisfies
289 ====== ================================
290 left ``self[i-1] < value <= self[i]``
291 right ``self[i-1] <= value < self[i]``
292 ====== ================================
294 Parameters
295 ----------
296 value : array-like, list or scalar
297 Value(s) to insert into `self`.
298 side : {'left', 'right'}, optional
299 If 'left', the index of the first suitable location found is given.
300 If 'right', return the last such index. If there is no suitable
301 index, return either 0 or N (where N is the length of `self`).
302 sorter : 1-D array-like, optional
303 Optional array of integer indices that sort array a into ascending
304 order. They are typically the result of argsort.
306 Returns
307 -------
308 array of ints or int
309 If value is array-like, array of insertion points.
310 If value is scalar, a single integer.
312 See Also
313 --------
314 numpy.searchsorted : Similar method from NumPy.
315 """
316 npvalue = self._validate_setitem_value(value)
317 return self._ndarray.searchsorted(npvalue, side=side, sorter=sorter)
319 def shift(self, periods: int = 1, fill_value=None) -> Self:
320 """
321 Shift values by desired number.
323 Newly introduced missing values are filled with
324 ``self.dtype.na_value``.
326 Parameters
327 ----------
328 periods : int, default 1
329 The number of periods to shift. Negative values are allowed
330 for shifting backwards.
332 fill_value : object, optional
333 The scalar value to use for newly introduced missing values.
334 The default is ``self.dtype.na_value``.
336 Returns
337 -------
338 ExtensionArray
339 Shifted.
341 Notes
342 -----
343 If ``self`` is empty or ``periods`` is 0, a copy of ``self`` is
344 returned.
346 If ``periods > len(self)``, then an array of size
347 len(self) is returned, with all values filled with
348 ``self.dtype.na_value``.
349 """
350 # NB: shift is always along axis=self.ndim-1
351 fill_value = self._validate_scalar(fill_value)
352 new_values = shift(self._ndarray, periods, fill_value)
354 return self._from_backing_data(new_values)
356 def __setitem__(self, key, value) -> None:
357 if self._readonly:
358 raise ValueError("Cannot modify read-only array")
360 key = check_array_indexer(self, key)
361 value = self._validate_setitem_value(value)
362 try:
363 self._ndarray[key] = value
364 except TypeError as err:
365 if self._ndarray.dtype.kind == "c":
366 # GH#54761 numpy raises TypeError for cases like setting a
367 # 1d sequence into a complex scalar position ("only 0-dimensional
368 # arrays can be converted to Python scalars"); normalize to
369 # ValueError for consistency with other numeric dtypes.
370 raise ValueError(*err.args) from err
371 raise
373 def _validate_setitem_value(self, value):
374 return value
376 @overload
377 def __getitem__(self, key: ScalarIndexer) -> Any: ...
379 @overload
380 def __getitem__(
381 self,
382 key: SequenceIndexer | PositionalIndexerTuple,
383 ) -> Self: ...
385 def __getitem__(
386 self,
387 key: PositionalIndexer2D,
388 ) -> Self | Any:
389 if lib.is_integer(key):
390 # fast-path
391 result = self._ndarray[key]
392 if self.ndim == 1:
393 return self._box_func(result)
394 result = self._from_backing_data(result)
395 if getitem_returns_view(self, key):
396 result._readonly = self._readonly
397 return result
399 # error: Incompatible types in assignment (expression has type "ExtensionArray",
400 # variable has type "Union[int, slice, ndarray]")
401 key = extract_array(key, extract_numpy=True) # type: ignore[assignment]
402 key = check_array_indexer(self, key)
403 result = self._ndarray[key]
404 if lib.is_scalar(result):
405 return self._box_func(result)
407 result = self._from_backing_data(result)
408 if getitem_returns_view(self, key):
409 result._readonly = self._readonly
410 return result
412 def _pad_or_backfill(
413 self,
414 *,
415 method: FillnaOptions,
416 limit: int | None = None,
417 limit_area: Literal["inside", "outside"] | None = None,
418 copy: bool = True,
419 ) -> Self:
420 mask = self.isna()
421 if mask.any():
422 # (for now) when self.ndim == 2, we assume axis=0
423 func = missing.get_fill_func(method, ndim=self.ndim)
425 npvalues = self._ndarray.T
426 if copy:
427 npvalues = npvalues.copy()
428 elif self._readonly:
429 raise ValueError("Cannot modify read-only array")
430 func(npvalues, limit=limit, limit_area=limit_area, mask=mask.T)
431 npvalues = npvalues.T
433 if copy:
434 new_values = self._from_backing_data(npvalues)
435 else:
436 new_values = self
438 elif copy:
439 new_values = self.copy()
440 else:
441 new_values = self
442 return new_values
444 def fillna(self, value, limit: int | None = None, copy: bool = True) -> Self:
445 """
446 Fill NA/NaN values using the specified method.
448 Parameters
449 ----------
450 value : scalar, array-like
451 If a scalar value is passed it is used to fill all missing values.
452 Alternatively, an array-like 'value' can be given. It's expected
453 that the array-like have the same length as 'self'.
454 limit : int, default None
455 The maximum number of entries along the entire axis where NaNs will be
456 filled.
458 Returns
459 -------
460 ExtensionArray
461 With NA/NaN filled.
462 """
463 mask = self.isna()
464 if limit is not None and limit < len(self):
465 # mypy doesn't like that mask can be an EA which need not have `cumsum`
466 modify = mask.cumsum() > limit # type: ignore[union-attr]
467 if modify.any():
468 # Only copy mask if necessary
469 mask = mask.copy()
470 mask[modify] = False
471 # error: Argument 2 to "check_value_size" has incompatible type
472 # "ExtensionArray"; expected "ndarray"
473 value = missing.check_value_size(
474 value,
475 mask, # type: ignore[arg-type]
476 len(self),
477 )
479 if mask.any():
480 # fill with value
481 if copy:
482 new_values = self.copy()
483 else:
484 new_values = self[:]
485 new_values[mask] = value
486 else:
487 # We validate the fill_value even if there is nothing to fill
488 self._validate_setitem_value(value)
490 if not copy:
491 new_values = self[:]
492 else:
493 new_values = self.copy()
494 return new_values
496 # ------------------------------------------------------------------------
497 # Reductions
499 def _wrap_reduction_result(self, axis: AxisInt | None, result) -> Any:
500 if axis is None or self.ndim == 1:
501 return self._box_func(result)
502 return self._from_backing_data(result)
504 # ------------------------------------------------------------------------
505 # __array_function__ methods
507 def _putmask(self, mask: npt.NDArray[np.bool_], value) -> None:
508 """
509 Analogue to np.putmask(self, mask, value)
511 Parameters
512 ----------
513 mask : np.ndarray[bool]
514 value : scalar or listlike
516 Raises
517 ------
518 TypeError
519 If value cannot be cast to self.dtype.
520 """
521 if self._readonly:
522 raise ValueError("Cannot modify read-only array")
523 value = self._validate_setitem_value(value)
525 np.putmask(self._ndarray, mask, value)
527 def _where(self: Self, mask: npt.NDArray[np.bool_], value) -> Self:
528 """
529 Analogue to np.where(mask, self, value)
531 Parameters
532 ----------
533 mask : np.ndarray[bool]
534 value : scalar or listlike
536 Raises
537 ------
538 TypeError
539 If value cannot be cast to self.dtype.
540 """
541 value = self._validate_setitem_value(value)
543 res_values = np.where(mask, self._ndarray, value)
544 if res_values.dtype != self._ndarray.dtype:
545 raise AssertionError(
546 # GH#56410
547 "Something has gone wrong, please report a bug at "
548 "github.com/pandas-dev/pandas/"
549 )
550 return self._from_backing_data(res_values)
552 # ------------------------------------------------------------------------
553 # Index compat methods
555 def insert(self, loc: int, item) -> Self:
556 """
557 Make new ExtensionArray inserting new item at location. Follows
558 Python list.append semantics for negative values.
560 Parameters
561 ----------
562 loc : int
563 item : object
565 Returns
566 -------
567 type(self)
568 """
569 loc = validate_insert_loc(loc, len(self))
571 code = self._validate_scalar(item)
573 new_vals = np.concatenate(
574 (
575 self._ndarray[:loc],
576 np.asarray([code], dtype=self._ndarray.dtype),
577 self._ndarray[loc:],
578 )
579 )
580 return self._from_backing_data(new_vals)
582 # ------------------------------------------------------------------------
583 # Additional array methods
584 # These are not part of the EA API, but we implement them because
585 # pandas assumes they're there.
587 def value_counts(self, dropna: bool = True) -> Series:
588 """
589 Return a Series containing counts of unique values.
591 Parameters
592 ----------
593 dropna : bool, default True
594 Don't include counts of NA values.
596 Returns
597 -------
598 Series
599 """
600 if self.ndim != 1:
601 raise NotImplementedError
603 from pandas import (
604 Index,
605 Series,
606 )
608 if dropna:
609 # error: Invalid index type
610 values = self[~self.isna()]._ndarray # type: ignore[index]
611 else:
612 values = self._ndarray
614 result = value_counts(values, sort=False, dropna=dropna)
616 index_arr = self._from_backing_data(np.asarray(result.index._data))
617 index = Index(index_arr, name=result.index.name, copy=False)
618 return Series(result._values, index=index, name=result.name, copy=False)
620 def _quantile(
621 self,
622 qs: npt.NDArray[np.float64],
623 interpolation: str,
624 ) -> Self:
625 # TODO: disable for Categorical if not ordered?
627 mask = np.asarray(self.isna())
628 arr = self._ndarray
629 fill_value = self._internal_fill_value
631 res_values = quantile_with_mask(arr, mask, fill_value, qs, interpolation)
632 if res_values.dtype == self._ndarray.dtype:
633 return self._from_backing_data(res_values)
634 else:
635 # e.g. test_quantile_empty we are empty integer dtype and res_values
636 # has floating dtype
637 # TODO: technically __init__ isn't defined here.
638 # Should we raise NotImplementedError and handle this on NumpyEA?
639 return type(self)(res_values) # type: ignore[call-arg]
641 # ------------------------------------------------------------------------
642 # numpy-like methods
644 @classmethod
645 def _empty(cls, shape: Shape, dtype: ExtensionDtype) -> Self:
646 """
647 Analogous to np.empty(shape, dtype=dtype)
649 Parameters
650 ----------
651 shape : tuple[int]
652 dtype : ExtensionDtype
653 """
654 # The base implementation uses a naive approach to find the dtype
655 # for the backing ndarray
656 arr = cls._from_sequence([], dtype=dtype)
657 backing = np.empty(shape, dtype=arr._ndarray.dtype)
658 return arr._from_backing_data(backing)