1from __future__ import annotations
2
3from typing import (
4 TYPE_CHECKING,
5 Any,
6 Literal,
7 Self,
8 cast,
9)
10
11import numpy as np
12
13from pandas._libs import lib
14from pandas._libs.tslibs import is_supported_dtype
15from pandas.compat.numpy import function as nv
16from pandas.errors import LossySetitemError
17from pandas.util._decorators import set_module
18
19from pandas.core.dtypes.astype import (
20 astype_array,
21 astype_float_to_int_nansafe,
22 astype_is_view,
23)
24from pandas.core.dtypes.cast import (
25 construct_1d_object_array_from_listlike,
26 maybe_downcast_to_dtype,
27 np_can_hold_element,
28)
29from pandas.core.dtypes.common import pandas_dtype
30from pandas.core.dtypes.dtypes import NumpyEADtype
31from pandas.core.dtypes.missing import (
32 is_valid_na_for_dtype,
33 isna,
34)
35
36from pandas.core import (
37 arraylike,
38 missing,
39 nanops,
40 ops,
41)
42from pandas.core.arraylike import OpsMixin
43from pandas.core.arrays._mixins import NDArrayBackedExtensionArray
44from pandas.core.construction import ensure_wrapped_if_datetimelike
45from pandas.core.strings.object_array import ObjectStringArrayMixin
46
47if TYPE_CHECKING:
48 from collections.abc import Callable
49
50 from pandas._typing import (
51 ArrayLike,
52 AxisInt,
53 Dtype,
54 FillnaOptions,
55 InterpolateOptions,
56 NpDtype,
57 NumpySorter,
58 NumpyValueArrayLike,
59 Scalar,
60 TakeIndexer,
61 npt,
62 )
63
64 from pandas import Index
65 from pandas.arrays import StringArray
66 from pandas.core.arrays.base import ExtensionArray
67
68
69@set_module("pandas.arrays")
70class NumpyExtensionArray(
71 OpsMixin,
72 NDArrayBackedExtensionArray,
73 ObjectStringArrayMixin,
74):
75 """
76 A pandas ExtensionArray for NumPy data.
77
78 This is mostly for internal compatibility, and is not especially
79 useful on its own.
80
81 Parameters
82 ----------
83 values : ndarray
84 The NumPy ndarray to wrap. Must be 1-dimensional.
85 copy : bool, default False
86 Whether to copy `values`.
87
88 Attributes
89 ----------
90 None
91
92 Methods
93 -------
94 None
95
96 See Also
97 --------
98 array : Create an array.
99 Series.to_numpy : Convert a Series to a NumPy array.
100
101 Examples
102 --------
103 >>> pd.arrays.NumpyExtensionArray(np.array([0, 1, 2, 3]))
104 <NumpyExtensionArray>
105 [0, 1, 2, 3]
106 Length: 4, dtype: int64
107 """
108
109 # If you're wondering why pd.Series(cls) doesn't put the array in an
110 # ExtensionBlock, search for `ABCNumpyExtensionArray`. We check for
111 # that _typ to ensure that users don't unnecessarily use EAs inside
112 # pandas internals, which turns off things like block consolidation.
113 _typ = "npy_extension"
114 __array_priority__ = 1000
115 _ndarray: np.ndarray
116 _dtype: NumpyEADtype
117 _internal_fill_value = np.nan
118
119 # ------------------------------------------------------------------------
120 # Constructors
121
122 def __init__(
123 self, values: np.ndarray | NumpyExtensionArray, copy: bool = False
124 ) -> None:
125 if isinstance(values, type(self)):
126 values = values._ndarray
127 if not isinstance(values, np.ndarray):
128 raise ValueError(
129 f"'values' must be a NumPy array, not {type(values).__name__}"
130 )
131
132 if values.ndim == 0:
133 # Technically we support 2, but do not advertise that fact.
134 raise ValueError("NumpyExtensionArray must be 1-dimensional.")
135
136 if copy:
137 values = values.copy()
138
139 dtype = NumpyEADtype(values.dtype)
140 super().__init__(values, dtype)
141
142 @classmethod
143 def _from_sequence(
144 cls, scalars, *, dtype: Dtype | None = None, copy: bool = False
145 ) -> NumpyExtensionArray:
146 if isinstance(dtype, NumpyEADtype):
147 dtype = dtype._dtype
148 if dtype is not None:
149 dtype = np.dtype(dtype) # type: ignore[arg-type]
150
151 if dtype is not None and dtype.kind in "iu":
152 # GH#41724 - validate NaN before casting float -> int
153 result = np.asarray(scalars)
154 if np.issubdtype(result.dtype, np.floating):
155 result = astype_float_to_int_nansafe(result, dtype, copy=copy)
156 else:
157 result = np.asarray(scalars, dtype=dtype)
158 else:
159 result = np.asarray(scalars, dtype=dtype)
160 if (
161 result.ndim > 1
162 and not hasattr(scalars, "dtype")
163 and (dtype is None or dtype == object)
164 ):
165 # e.g. list-of-tuples
166 result = construct_1d_object_array_from_listlike(scalars)
167
168 if copy and result is scalars:
169 result = result.copy()
170 return cls(result)
171
172 def _validate_setitem_value(self, value):
173 if isinstance(value, type(self)):
174 value = value._ndarray
175
176 # Match Block._standardize_fill_value behavior
177 if self._ndarray.dtype.kind != "O" and is_valid_na_for_dtype(
178 value, self._ndarray.dtype
179 ):
180 value = self.dtype.na_value
181
182 try:
183 return np_can_hold_element(self._ndarray.dtype, value)
184 except LossySetitemError as err:
185 raise TypeError(
186 f"Invalid value '{value!s}' for dtype '{self.dtype}'"
187 ) from err
188 except NotImplementedError:
189 # np_can_hold_element doesn't handle all dtypes (e.g. "U"),
190 # fall back to no validation for those.
191 return value
192
193 def searchsorted(
194 self,
195 value: NumpyValueArrayLike | ExtensionArray,
196 side: Literal["left", "right"] = "left",
197 sorter: NumpySorter | None = None,
198 ) -> npt.NDArray[np.intp] | np.intp:
199 # Parent's searchsorted calls _validate_setitem_value, which is
200 # too strict for search (e.g. rejects float into int). Delegate
201 # directly to numpy which handles cross-dtype searches correctly.
202 return self._ndarray.searchsorted(value, side=side, sorter=sorter) # type: ignore[arg-type]
203
204 def _cast_pointwise_result(self, values) -> ArrayLike:
205 result = super()._cast_pointwise_result(values)
206 lkind = self.dtype.kind
207 rkind = result.dtype.kind
208 if (
209 (lkind in "iu" and rkind in "iu")
210 or (lkind == "f" and rkind == "f")
211 or (lkind == rkind == "c")
212 ):
213 result = maybe_downcast_to_dtype(result, self.dtype.numpy_dtype)
214 elif rkind in "Mm":
215 # Ensure potential subsequent .astype(object) doesn't incorrectly
216 # convert Timestamps/Timedeltas to ints
217 from pandas import array as pd_array
218
219 result = pd_array(result, copy=False)
220 if isinstance(result, np.ndarray) and isinstance(self.dtype, NumpyEADtype):
221 # Subclasses like StringArray have their own validation;
222 # only wrap for plain NumpyEADtype.
223 result = type(self)(result)
224 return result
225
226 # ------------------------------------------------------------------------
227 # Data
228
229 @property
230 def dtype(self) -> NumpyEADtype:
231 return self._dtype
232
233 # ------------------------------------------------------------------------
234 # NumPy Array Interface
235
236 def __array__(
237 self, dtype: np.dtype | None = None, copy: bool | None = None
238 ) -> np.ndarray:
239 if copy is not None:
240 # Note: branch avoids `copy=None` for NumPy 1.x support
241 result = np.array(self._ndarray, dtype=dtype, copy=copy)
242 else:
243 result = np.asarray(self._ndarray, dtype=dtype)
244
245 if (
246 self._readonly
247 and not copy
248 and (dtype is None or astype_is_view(self.dtype, dtype))
249 ):
250 result = result.view()
251 result.flags.writeable = False
252
253 return result
254
255 def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
256 # Lightly modified version of
257 # https://numpy.org/doc/stable/reference/generated/numpy.lib.mixins.NDArrayOperatorsMixin.html
258 # The primary modification is not boxing scalar return values
259 # in NumpyExtensionArray, since pandas' ExtensionArrays are 1-d.
260
261 # this path never reaches ExtensionArray.__array_ufunc__
262 ops.disallow_datetimelike_logical_ufunc(ufunc, inputs)
263
264 out = kwargs.get("out", ())
265
266 result = arraylike.maybe_dispatch_ufunc_to_dunder_op(
267 self, ufunc, method, *inputs, **kwargs
268 )
269 if result is not NotImplemented:
270 return result
271
272 if "out" in kwargs:
273 # e.g. test_ufunc_unary
274 return arraylike.dispatch_ufunc_with_out(
275 self, ufunc, method, *inputs, **kwargs
276 )
277
278 if method == "reduce":
279 result = arraylike.dispatch_reduction_ufunc(
280 self, ufunc, method, *inputs, **kwargs
281 )
282 if result is not NotImplemented:
283 # e.g. tests.series.test_ufunc.TestNumpyReductions
284 return result
285
286 # Defer to the implementation of the ufunc on unwrapped values.
287 inputs = tuple(
288 x._ndarray if isinstance(x, NumpyExtensionArray) else x for x in inputs
289 )
290 if out:
291 kwargs["out"] = tuple(
292 x._ndarray if isinstance(x, NumpyExtensionArray) else x for x in out
293 )
294 result = getattr(ufunc, method)(*inputs, **kwargs)
295
296 if ufunc.nout > 1:
297 # multiple return values; re-box array-like results
298 return tuple(type(self)(x) for x in result)
299 elif method == "at":
300 # no return value
301 return None
302 elif method == "reduce":
303 if isinstance(result, np.ndarray):
304 # e.g. test_np_reduce_2d
305 return type(self)(result)
306
307 # e.g. test_np_max_nested_tuples
308 return result
309 else:
310 if self.dtype.type is str: # type: ignore[comparison-overlap]
311 # StringDtype
312 self = cast("StringArray", self)
313 try:
314 # specify dtype to preserve storage/na_value
315 return type(self)(result, dtype=self.dtype)
316 except ValueError:
317 # if validation of input fails (no strings)
318 # -> fallback to returning raw numpy array
319 return result
320 # one return value; re-box array-like results
321 return type(self)(result)
322
323 # ------------------------------------------------------------------------
324 # Pandas ExtensionArray Interface
325
326 def astype(self, dtype, copy: bool = True):
327 dtype = pandas_dtype(dtype)
328
329 if dtype == self.dtype:
330 if copy:
331 return self.copy()
332 return self
333
334 result = astype_array(self._ndarray, dtype=dtype, copy=copy)
335 return result
336
337 def isna(self) -> np.ndarray:
338 return isna(self._ndarray)
339
340 def _validate_scalar(self, fill_value):
341 if fill_value is None:
342 # Primarily for subclasses
343 fill_value = self.dtype.na_value
344 return fill_value
345
346 def _values_for_factorize(self) -> tuple[np.ndarray, float | None]:
347 if self.dtype.kind in "iub":
348 fv = None
349 else:
350 fv = np.nan
351 return self._ndarray, fv
352
353 # Base EA class (and all other EA classes) don't have limit_area keyword
354 # This can be removed here as well when the interpolate ffill/bfill method
355 # deprecation is enforced
356 def _pad_or_backfill(
357 self,
358 *,
359 method: FillnaOptions,
360 limit: int | None = None,
361 limit_area: Literal["inside", "outside"] | None = None,
362 copy: bool = True,
363 ) -> Self:
364 """
365 ffill or bfill along axis=0.
366 """
367 if copy:
368 out_data = self._ndarray.copy()
369 else:
370 out_data = self._ndarray
371
372 meth = missing.clean_fill_method(method)
373 missing.pad_or_backfill_inplace(
374 out_data.T,
375 method=meth,
376 axis=0,
377 limit=limit,
378 limit_area=limit_area,
379 )
380
381 if not copy:
382 return self
383 return type(self)._simple_new(out_data, dtype=self.dtype)
384
385 def interpolate(
386 self,
387 *,
388 method: InterpolateOptions,
389 axis: int,
390 index: Index,
391 limit,
392 limit_direction,
393 limit_area,
394 copy: bool,
395 **kwargs,
396 ) -> Self:
397 """
398 See NDFrame.interpolate.__doc__.
399 """
400 # NB: we return type(self) even if copy=False
401 if not self.dtype._is_numeric:
402 raise TypeError(f"Cannot interpolate with {self.dtype} dtype")
403
404 if not copy:
405 out_data = self._ndarray
406 else:
407 out_data = self._ndarray.copy()
408
409 # TODO: assert we have floating dtype?
410 missing.interpolate_2d_inplace(
411 out_data,
412 method=method,
413 axis=axis,
414 index=index,
415 limit=limit,
416 limit_direction=limit_direction,
417 limit_area=limit_area,
418 **kwargs,
419 )
420 if not copy:
421 return self
422 return type(self)._simple_new(out_data, dtype=self.dtype)
423
424 def take(
425 self,
426 indices: TakeIndexer,
427 *,
428 allow_fill: bool = False,
429 fill_value: Any = None,
430 axis: AxisInt = 0,
431 ) -> Self:
432 """
433 Take entries from this array at each index in a list of indices,
434 producing an array containing only those entries.
435 """
436 result = super().take(
437 indices, allow_fill=allow_fill, fill_value=fill_value, axis=axis
438 )
439 # See GH#62448.
440 if self.dtype.kind in "iub":
441 return type(self)(result._ndarray, copy=False)
442
443 return result
444
445 # ------------------------------------------------------------------------
446 # Reductions
447
448 def any(
449 self,
450 *,
451 axis: AxisInt | None = None,
452 out=None,
453 keepdims: bool = False,
454 skipna: bool = True,
455 ):
456 nv.validate_any((), {"out": out, "keepdims": keepdims})
457 result = nanops.nanany(self._ndarray, axis=axis, skipna=skipna)
458 return self._wrap_reduction_result(axis, result)
459
460 def all(
461 self,
462 *,
463 axis: AxisInt | None = None,
464 out=None,
465 keepdims: bool = False,
466 skipna: bool = True,
467 ):
468 nv.validate_all((), {"out": out, "keepdims": keepdims})
469 result = nanops.nanall(self._ndarray, axis=axis, skipna=skipna)
470 return self._wrap_reduction_result(axis, result)
471
472 def min(
473 self, *, axis: AxisInt | None = None, skipna: bool = True, **kwargs
474 ) -> Scalar:
475 nv.validate_min((), kwargs)
476 result = nanops.nanmin(
477 values=self._ndarray, axis=axis, mask=self.isna(), skipna=skipna
478 )
479 return self._wrap_reduction_result(axis, result)
480
481 def max(
482 self, *, axis: AxisInt | None = None, skipna: bool = True, **kwargs
483 ) -> Scalar:
484 nv.validate_max((), kwargs)
485 result = nanops.nanmax(
486 values=self._ndarray, axis=axis, mask=self.isna(), skipna=skipna
487 )
488 return self._wrap_reduction_result(axis, result)
489
490 def sum(
491 self,
492 *,
493 axis: AxisInt | None = None,
494 skipna: bool = True,
495 min_count: int = 0,
496 **kwargs,
497 ) -> Scalar:
498 nv.validate_sum((), kwargs)
499 result = nanops.nansum(
500 self._ndarray, axis=axis, skipna=skipna, min_count=min_count
501 )
502 return self._wrap_reduction_result(axis, result)
503
504 def prod(
505 self,
506 *,
507 axis: AxisInt | None = None,
508 skipna: bool = True,
509 min_count: int = 0,
510 **kwargs,
511 ) -> Scalar:
512 nv.validate_prod((), kwargs)
513 result = nanops.nanprod(
514 self._ndarray, axis=axis, skipna=skipna, min_count=min_count
515 )
516 return self._wrap_reduction_result(axis, result)
517
518 def mean(
519 self,
520 *,
521 axis: AxisInt | None = None,
522 dtype: NpDtype | None = None,
523 out=None,
524 keepdims: bool = False,
525 skipna: bool = True,
526 ):
527 nv.validate_mean((), {"dtype": dtype, "out": out, "keepdims": keepdims})
528 result = nanops.nanmean(self._ndarray, axis=axis, skipna=skipna)
529 return self._wrap_reduction_result(axis, result)
530
531 def median(
532 self,
533 *,
534 axis: AxisInt | None = None,
535 out=None,
536 overwrite_input: bool = False,
537 keepdims: bool = False,
538 skipna: bool = True,
539 ):
540 nv.validate_median(
541 (), {"out": out, "overwrite_input": overwrite_input, "keepdims": keepdims}
542 )
543 result = nanops.nanmedian(self._ndarray, axis=axis, skipna=skipna)
544 return self._wrap_reduction_result(axis, result)
545
546 def std(
547 self,
548 *,
549 axis: AxisInt | None = None,
550 dtype: NpDtype | None = None,
551 out=None,
552 ddof: int = 1,
553 keepdims: bool = False,
554 skipna: bool = True,
555 ):
556 nv.validate_stat_ddof_func(
557 (), {"dtype": dtype, "out": out, "keepdims": keepdims}, fname="std"
558 )
559 result = nanops.nanstd(self._ndarray, axis=axis, skipna=skipna, ddof=ddof)
560 return self._wrap_reduction_result(axis, result)
561
562 def var(
563 self,
564 *,
565 axis: AxisInt | None = None,
566 dtype: NpDtype | None = None,
567 out=None,
568 ddof: int = 1,
569 keepdims: bool = False,
570 skipna: bool = True,
571 ):
572 nv.validate_stat_ddof_func(
573 (), {"dtype": dtype, "out": out, "keepdims": keepdims}, fname="var"
574 )
575 result = nanops.nanvar(self._ndarray, axis=axis, skipna=skipna, ddof=ddof)
576 return self._wrap_reduction_result(axis, result)
577
578 def sem(
579 self,
580 *,
581 axis: AxisInt | None = None,
582 dtype: NpDtype | None = None,
583 out=None,
584 ddof: int = 1,
585 keepdims: bool = False,
586 skipna: bool = True,
587 ):
588 nv.validate_stat_ddof_func(
589 (), {"dtype": dtype, "out": out, "keepdims": keepdims}, fname="sem"
590 )
591 result = nanops.nansem(self._ndarray, axis=axis, skipna=skipna, ddof=ddof)
592 return self._wrap_reduction_result(axis, result)
593
594 def kurt(
595 self,
596 *,
597 axis: AxisInt | None = None,
598 dtype: NpDtype | None = None,
599 out=None,
600 keepdims: bool = False,
601 skipna: bool = True,
602 ):
603 nv.validate_stat_ddof_func(
604 (), {"dtype": dtype, "out": out, "keepdims": keepdims}, fname="kurt"
605 )
606 result = nanops.nankurt(self._ndarray, axis=axis, skipna=skipna)
607 return self._wrap_reduction_result(axis, result)
608
609 def skew(
610 self,
611 *,
612 axis: AxisInt | None = None,
613 dtype: NpDtype | None = None,
614 out=None,
615 keepdims: bool = False,
616 skipna: bool = True,
617 ):
618 nv.validate_stat_ddof_func(
619 (), {"dtype": dtype, "out": out, "keepdims": keepdims}, fname="skew"
620 )
621 result = nanops.nanskew(self._ndarray, axis=axis, skipna=skipna)
622 return self._wrap_reduction_result(axis, result)
623
624 # ------------------------------------------------------------------------
625 # Additional Methods
626
627 def to_numpy(
628 self,
629 dtype: npt.DTypeLike | None = None,
630 copy: bool = False,
631 na_value: object = lib.no_default,
632 ) -> np.ndarray:
633 if na_value is not lib.no_default and (mask := self.isna()).any():
634 result = self._ndarray.copy()
635 result[mask] = na_value
636 else:
637 result = self._ndarray
638 if not copy and self._readonly:
639 result = result.view()
640 result.flags.writeable = False
641
642 result = np.asarray(result, dtype=dtype)
643
644 if copy and result is self._ndarray:
645 result = result.copy()
646
647 return result
648
649 # ------------------------------------------------------------------------
650 # Ops
651
652 def __invert__(self) -> NumpyExtensionArray:
653 return type(self)(~self._ndarray)
654
655 def __neg__(self) -> NumpyExtensionArray:
656 return type(self)(-self._ndarray)
657
658 def __pos__(self) -> NumpyExtensionArray:
659 return type(self)(+self._ndarray)
660
661 def __abs__(self) -> NumpyExtensionArray:
662 return type(self)(abs(self._ndarray))
663
664 def _cmp_method(self, other, op):
665 if isinstance(other, NumpyExtensionArray):
666 other = other._ndarray
667
668 other = ops.maybe_prepare_scalar_for_op(other, (len(self),))
669 pd_op = ops.get_array_op(op)
670 other = ensure_wrapped_if_datetimelike(other)
671 result = pd_op(self._ndarray, other)
672
673 if op is divmod or op is ops.rdivmod:
674 a, b = result
675 if isinstance(a, np.ndarray):
676 # for e.g. op vs TimedeltaArray, we may already
677 # have an ExtensionArray, in which case we do not wrap
678 return self._wrap_ndarray_result(a), self._wrap_ndarray_result(b)
679 return a, b
680
681 if isinstance(result, np.ndarray):
682 # for e.g. multiplication vs TimedeltaArray, we may already
683 # have an ExtensionArray, in which case we do not wrap
684 return self._wrap_ndarray_result(result)
685 return result
686
687 _arith_method = _cmp_method
688
689 def _wrap_ndarray_result(self, result: np.ndarray):
690 # If we have timedelta64[ns] result, return a TimedeltaArray instead
691 # of a NumpyExtensionArray
692 if result.dtype.kind == "m" and is_supported_dtype(result.dtype):
693 from pandas.core.arrays import TimedeltaArray
694
695 return TimedeltaArray._simple_new(result, dtype=result.dtype)
696 return type(self)(result)
697
698 def _formatter(self, boxed: bool = False) -> Callable[[Any], str | None]:
699 # NEP 51: https://github.com/numpy/numpy/pull/22449
700 if self.dtype.kind == "U":
701 # not "S": str(np.bytes_(b"foo")) is already "b'foo'"
702 return "'{}'".format
703 elif self.dtype == "object":
704 return repr
705 else:
706 return str