1from __future__ import annotations
2
3from typing import (
4 TYPE_CHECKING,
5 Any,
6 Literal,
7 Self,
8 cast,
9 overload,
10)
11import warnings
12
13import numpy as np
14
15from pandas._config import (
16 is_nan_na,
17 using_python_scalars,
18)
19
20from pandas._libs import (
21 algos as libalgos,
22 groupby as libgroupby,
23 lib,
24 missing as libmissing,
25)
26from pandas._libs.tslibs import is_supported_dtype
27from pandas.compat import (
28 IS64,
29 is_platform_windows,
30)
31from pandas.errors import (
32 AbstractMethodError,
33 Pandas4Warning,
34)
35from pandas.util._exceptions import find_stack_level
36
37from pandas.core.dtypes.astype import astype_is_view
38from pandas.core.dtypes.base import ExtensionDtype
39from pandas.core.dtypes.cast import (
40 construct_1d_object_array_from_listlike,
41 maybe_downcast_to_dtype,
42)
43from pandas.core.dtypes.common import (
44 is_bool,
45 is_float_dtype,
46 is_integer_dtype,
47 is_list_like,
48 is_scalar,
49 is_string_dtype,
50 pandas_dtype,
51)
52from pandas.core.dtypes.dtypes import (
53 ArrowDtype,
54 BaseMaskedDtype,
55)
56from pandas.core.dtypes.missing import (
57 array_equivalent,
58 is_valid_na_for_dtype,
59 isna,
60 notna,
61)
62
63from pandas.core import (
64 algorithms as algos,
65 arraylike,
66 missing,
67 nanops,
68 ops,
69)
70from pandas.core.algorithms import (
71 factorize_array,
72 isin,
73 mode,
74 take,
75)
76from pandas.core.array_algos import (
77 masked_accumulations,
78 masked_reductions,
79)
80from pandas.core.array_algos.quantile import quantile_with_mask
81from pandas.core.array_algos.transforms import shift
82from pandas.core.arraylike import OpsMixin
83from pandas.core.arrays._utils import to_numpy_dtype_inference
84from pandas.core.arrays.base import ExtensionArray
85from pandas.core.construction import (
86 array as pd_array,
87 ensure_wrapped_if_datetimelike,
88 extract_array,
89)
90from pandas.core.indexers import (
91 check_array_indexer,
92 getitem_returns_view,
93)
94from pandas.core.ops import invalid_comparison
95from pandas.core.util.hashing import hash_array
96
97if TYPE_CHECKING:
98 from collections.abc import Callable
99 from collections.abc import (
100 Iterator,
101 Sequence,
102 )
103 from pandas import Series
104 from pandas.core.arrays import BooleanArray
105 from pandas._typing import (
106 Dtype,
107 NumpySorter,
108 NumpyValueArrayLike,
109 ArrayLike,
110 AstypeArg,
111 AxisInt,
112 DtypeObj,
113 FillnaOptions,
114 InterpolateOptions,
115 NpDtype,
116 PositionalIndexer,
117 RankMethod,
118 RankNaOption,
119 Scalar,
120 ScalarIndexer,
121 SequenceIndexer,
122 Shape,
123 npt,
124 )
125 from pandas._libs.missing import NAType
126 from pandas.core.arrays import FloatingArray
127
128from pandas.compat.numpy import function as nv
129
130
131class BaseMaskedArray(OpsMixin, ExtensionArray):
132 """
133 Base class for masked arrays (which use _data and _mask to store the data).
134
135 numpy based
136 """
137
138 # our underlying data and mask are each ndarrays
139 _data: np.ndarray
140 _mask: npt.NDArray[np.bool_]
141
142 @classmethod
143 def _simple_new(cls, values: np.ndarray, mask: npt.NDArray[np.bool_]) -> Self:
144 result = BaseMaskedArray.__new__(cls)
145 result._data = values
146 result._mask = mask
147 return result
148
149 def __init__(
150 self, values: np.ndarray, mask: npt.NDArray[np.bool_], copy: bool = False
151 ) -> None:
152 # values is supposed to already be validated in the subclass
153 if not (isinstance(mask, np.ndarray) and mask.dtype == np.bool_):
154 raise TypeError(
155 "mask should be boolean numpy array. Use "
156 "the 'pd.array' function instead"
157 )
158 if values.shape != mask.shape:
159 raise ValueError("values.shape must match mask.shape")
160
161 if copy:
162 values = values.copy()
163 mask = mask.copy()
164
165 self._data = values
166 self._mask = mask
167
168 @classmethod
169 def _from_sequence(cls, scalars, *, dtype=None, copy: bool = False) -> Self:
170 values, mask = cls._coerce_to_array(scalars, dtype=dtype, copy=copy)
171 return cls(values, mask)
172
173 def _cast_pointwise_result(self, values) -> ArrayLike:
174 if isna(values).all():
175 try:
176 return type(self)._from_sequence(values, dtype=self.dtype)
177 except TypeError:
178 # e.g. NaT goes through inference instead, GH#70233
179 pass
180 if not (isinstance(values, np.ndarray) and values.dtype == object):
181 values = construct_1d_object_array_from_listlike(values)
182 result = lib.maybe_convert_objects(values, convert_to_nullable_dtype=True)
183 lkind = self.dtype.kind
184 rkind = result.dtype.kind
185 if (lkind in "iu" and rkind in "iu") or (lkind == rkind == "f"):
186 result = cast("BaseMaskedArray", result)
187 new_data = maybe_downcast_to_dtype(
188 result._data, dtype=self.dtype.numpy_dtype
189 )
190 result = type(result)(new_data, result._mask)
191 return result
192
193 @classmethod
194 def _empty(cls, shape: Shape, dtype: ExtensionDtype) -> Self:
195 """
196 Create an ExtensionArray with the given shape and dtype.
197
198 See also
199 --------
200 ExtensionDtype.empty
201 ExtensionDtype.empty is the 'official' public version of this API.
202 """
203 dtype = cast("BaseMaskedDtype", dtype)
204 values: np.ndarray = np.empty(shape, dtype=dtype.type)
205 values.fill(dtype._internal_fill_value)
206 mask = np.ones(shape, dtype=bool)
207 result = cls(values, mask)
208 if not isinstance(result, cls) or dtype != result.dtype:
209 raise NotImplementedError(
210 f"Default 'empty' implementation is invalid for dtype='{dtype}'"
211 )
212 return result
213
214 def _formatter(self, boxed: bool = False) -> Callable[[Any], str | None]:
215 # NEP 51: https://github.com/numpy/numpy/pull/22449
216 return str
217
218 @property
219 def dtype(self) -> BaseMaskedDtype:
220 raise AbstractMethodError(self)
221
222 @overload
223 def __getitem__(self, item: ScalarIndexer) -> Any: ...
224
225 @overload
226 def __getitem__(self, item: SequenceIndexer) -> Self: ...
227
228 def __getitem__(self, item: PositionalIndexer) -> Self | Any:
229 item = check_array_indexer(self, item)
230
231 newmask = self._mask[item]
232 if is_bool(newmask):
233 # This is a scalar indexing
234 if newmask:
235 return self.dtype.na_value
236 return self._data[item]
237
238 result = self._simple_new(self._data[item], newmask)
239 if getitem_returns_view(self, item):
240 result._readonly = self._readonly
241 return result
242
243 def _pad_or_backfill(
244 self,
245 *,
246 method: FillnaOptions,
247 limit: int | None = None,
248 limit_area: Literal["inside", "outside"] | None = None,
249 copy: bool = True,
250 ) -> Self:
251 mask = self._mask
252
253 if mask.any():
254 func = missing.get_fill_func(method, ndim=self.ndim)
255
256 npvalues = self._data.T
257 new_mask = mask.T
258 if copy:
259 npvalues = npvalues.copy()
260 new_mask = new_mask.copy()
261 else:
262 if self._readonly:
263 raise ValueError("Cannot modify read-only array")
264 if limit_area is not None:
265 mask = mask.copy()
266 func(npvalues, limit=limit, mask=new_mask)
267
268 if limit_area is not None and not mask.all():
269 mask = mask.T
270 neg_mask = ~mask
271 first = neg_mask.argmax()
272 last = len(neg_mask) - neg_mask[::-1].argmax() - 1
273 if limit_area == "inside":
274 new_mask[:first] |= mask[:first]
275 new_mask[last + 1 :] |= mask[last + 1 :]
276 elif limit_area == "outside":
277 new_mask[first + 1 : last] |= mask[first + 1 : last]
278
279 if copy:
280 return self._simple_new(npvalues.T, new_mask.T)
281 else:
282 return self
283 elif copy:
284 new_values = self.copy()
285 else:
286 new_values = self
287 return new_values
288
289 def fillna(self, value, limit: int | None = None, copy: bool = True) -> Self:
290 """
291 Fill NA/NaN values using the specified method.
292
293 Parameters
294 ----------
295 value : scalar, array-like
296 If a scalar value is passed it is used to fill all missing values.
297 Alternatively, an array-like "value" can be given. It's expected
298 that the array-like have the same length as 'self'.
299 limit : int, default None
300 The maximum number of entries where NA values will be filled.
301 copy : bool, default True
302 Whether to make a copy of the data before filling. If False, then
303 the original should be modified and no new memory should be allocated.
304 For ExtensionArray subclasses that cannot do this, it is at the
305 author's discretion whether to ignore "copy=False" or to raise.
306
307 Returns
308 -------
309 ExtensionArray
310 With NA/NaN filled.
311
312 See Also
313 --------
314 api.extensions.ExtensionArray.dropna : Return ExtensionArray without
315 NA values.
316 api.extensions.ExtensionArray.isna : A 1-D array indicating if
317 each value is missing.
318
319 Examples
320 --------
321 >>> arr = pd.array([np.nan, np.nan, 2, 3, np.nan, np.nan])
322 >>> arr.fillna(0)
323 <IntegerArray>
324 [0, 0, 2, 3, 0, 0]
325 Length: 6, dtype: Int64
326 """
327 mask = self._mask
328
329 # Fast path: single-pass Cython for scalar fill values. GH#42147
330 # Skip when fill value is itself NA — filling NAs with NA is a no-op,
331 # and the Cython path would incorrectly clear mask bits.
332 # When NaN and NA are distinguished (is_nan_na() is False), float NaN
333 # is a valid fill value, not NA.
334 value_is_na = is_valid_na_for_dtype(value, self.dtype) and not (
335 lib.is_float(value) and not is_nan_na()
336 )
337 if lib.is_scalar(value) and not value_is_na:
338 if not mask.any():
339 return self.copy() if copy else self[:]
340 try:
341 casted = self._validate_setitem_value(value)
342 except TypeError:
343 pass
344 else:
345 if copy:
346 new_data = self._data.copy()
347 new_mask = mask.copy()
348 else:
349 if self._readonly:
350 raise ValueError("Cannot modify read-only array")
351 new_data = self._data
352 new_mask = mask
353
354 if new_data.dtype.kind == "b":
355 # bool arrays need uint8 view for Cython memoryview
356 libalgos.scalar_fillna_inplace(
357 new_data.view("u1"),
358 np.uint8(casted),
359 mask=new_mask.view("u1"),
360 limit=limit,
361 )
362 else:
363 libalgos.scalar_fillna_inplace(
364 new_data,
365 casted,
366 mask=new_mask.view("u1"),
367 limit=limit,
368 )
369
370 return type(self)(new_data, new_mask, copy=False)
371
372 if limit is not None and limit < len(self):
373 modify = mask.cumsum() > limit
374 if modify.any():
375 # Only copy mask if necessary
376 mask = mask.copy()
377 mask[modify] = False
378
379 value = missing.check_value_size(value, mask, len(self))
380
381 if mask.any():
382 # fill with value
383 if copy:
384 new_values = self.copy()
385 else:
386 new_values = self[:]
387 new_values[mask] = value
388 elif copy:
389 new_values = self.copy()
390 else:
391 new_values = self[:]
392 return new_values
393
394 @classmethod
395 def _coerce_to_array(
396 cls, values, *, dtype: DtypeObj, copy: bool = False
397 ) -> tuple[np.ndarray, np.ndarray]:
398 raise AbstractMethodError(cls)
399
400 def _validate_setitem_value(self, value):
401 """
402 Check if we have a scalar that we can cast losslessly.
403
404 Raises
405 ------
406 TypeError
407 """
408 kind = self.dtype.kind
409 # TODO: get this all from np_can_hold_element?
410 if kind == "b":
411 if lib.is_bool(value):
412 return value
413
414 elif kind == "f":
415 if lib.is_integer(value) or lib.is_float(value):
416 return value
417
418 elif lib.is_integer(value) or (lib.is_float(value) and value.is_integer()):
419 return value
420 # TODO: unsigned checks
421
422 # Note: without the "str" here, the f-string rendering raises in
423 # py38 builds.
424 raise TypeError(f"Invalid value '{value!s}' for dtype '{self.dtype}'")
425
426 def insert(self, loc: int, item) -> Self:
427 if not is_valid_na_for_dtype(item, self.dtype):
428 self._validate_setitem_value(item)
429 return super().insert(loc, item)
430
431 def _validate_listlike(self, value) -> tuple[np.ndarray, npt.NDArray[np.bool_]]:
432 """
433 Validate a non-scalar setitem value and return ``(data, mask)``.
434
435 Raises
436 ------
437 TypeError
438 If `value` cannot be losslessly stored in self.dtype.
439 """
440 kind = self.dtype.kind
441
442 if hasattr(value, "dtype"):
443 their_kind = value.dtype.kind
444 # Compatible numeric/bool arrays (ndarray or masked) defer to
445 # _coerce_to_array, which handles precision/NaN checks via
446 # _safe_cast. Bool values are accepted for numeric targets to
447 # preserve the long-standing bool-as-int treatment exercised by
448 # Series.mask.
449 if (kind == "b" and their_kind == "b") or (
450 kind in "iuf" and their_kind in "iufb"
451 ):
452 return self._coerce_to_array(value, dtype=self.dtype)
453 elif not is_list_like(value):
454 # is_scalar in __setitem__ misses some non-listlike inputs
455 # (e.g. an opaque ``object()`` instance); route them through
456 # the scalar validator, which will raise.
457 self._validate_setitem_value(value)
458 else:
459 # Materialize other list-likes (e.g. a tuple, for which isna
460 # returns a scalar) as an object ndarray so that
461 # _cast_pointwise_result infers their type element-wise.
462 value = np.asarray(value, dtype=object)
463
464 # Otherwise let _cast_pointwise_result infer the natural type of
465 # `value` (without the silent string-/bool-to-numeric coercions
466 # that _coerce_to_array performs). A kind-mismatch is invalid.
467 casted = self._cast_pointwise_result(value)
468 casted_kind = casted.dtype.kind
469 if not (
470 (kind == "b" and casted_kind == "b")
471 or (kind in "iuf" and casted_kind in "iuf")
472 ):
473 raise TypeError(f"Invalid value '{value!s}' for dtype '{self.dtype}'")
474 return self._coerce_to_array(casted, dtype=self.dtype)
475
476 def __setitem__(self, key, value) -> None:
477 if self._readonly:
478 raise ValueError("Cannot modify read-only array")
479
480 key = check_array_indexer(self, key)
481
482 if is_scalar(value):
483 if is_valid_na_for_dtype(value, self.dtype) and not (
484 lib.is_float(value) and not is_nan_na()
485 ):
486 self._mask[key] = True
487 else:
488 value = self._validate_setitem_value(value)
489 self._data[key] = value
490 self._mask[key] = False
491 return
492
493 value, mask = self._validate_listlike(value)
494
495 self._data[key] = value
496 self._mask[key] = mask
497
498 def __contains__(self, key) -> bool:
499 if isna(key) and key is not self.dtype.na_value:
500 # GH#52840
501 if lib.is_float(key) and is_nan_na():
502 key = self.dtype.na_value
503 elif self._data.dtype.kind == "f" and lib.is_float(key):
504 return bool((np.isnan(self._data) & ~self._mask).any())
505
506 return bool(super().__contains__(key))
507
508 def __iter__(self) -> Iterator[Any]:
509 if self.ndim == 1:
510 if not self._hasna:
511 for val in self._data:
512 yield val
513 else:
514 na_value = self.dtype.na_value
515 for isna_, val in zip(self._mask, self._data, strict=True):
516 if isna_:
517 yield na_value
518 else:
519 yield val
520 else:
521 for i in range(len(self)):
522 yield self[i]
523
524 def __len__(self) -> int:
525 return len(self._data)
526
527 @property
528 def shape(self) -> Shape:
529 return self._data.shape
530
531 @property
532 def ndim(self) -> int:
533 return self._data.ndim
534
535 def swapaxes(self, axis1, axis2) -> Self:
536 data = self._data.swapaxes(axis1, axis2)
537 mask = self._mask.swapaxes(axis1, axis2)
538 return self._simple_new(data, mask)
539
540 def delete(self, loc, axis: AxisInt = 0) -> Self:
541 data = np.delete(self._data, loc, axis=axis)
542 mask = np.delete(self._mask, loc, axis=axis)
543 return self._simple_new(data, mask)
544
545 def reshape(self, *args, **kwargs) -> Self:
546 data = self._data.reshape(*args, **kwargs)
547 mask = self._mask.reshape(*args, **kwargs)
548 return self._simple_new(data, mask)
549
550 def ravel(self, *args, **kwargs) -> Self:
551 # TODO: need to make sure we have the same order for data/mask
552 data = self._data.ravel(*args, **kwargs)
553 mask = self._mask.ravel(*args, **kwargs)
554 return type(self)(data, mask)
555
556 def shift(self, periods: int = 1, fill_value=None) -> Self:
557 # NB: shift is always along axis=self.ndim-1
558 if fill_value is None:
559 new_data = shift(self._data, periods, 0)
560 new_mask = shift(self._mask, periods, True)
561 else:
562 new_data = shift(self._data, periods, fill_value)
563 new_mask = shift(self._mask, periods, False)
564 return type(self)(new_data, new_mask)
565
566 @property
567 def T(self) -> Self:
568 return self._simple_new(self._data.T, self._mask.T)
569
570 def round(self, decimals: int = 0, *args, **kwargs):
571 """
572 Round each value in the array a to the given number of decimals.
573
574 Parameters
575 ----------
576 decimals : int, default 0
577 Number of decimal places to round to. If decimals is negative,
578 it specifies the number of positions to the left of the decimal point.
579 *args, **kwargs
580 Additional arguments and keywords have no effect but might be
581 accepted for compatibility with NumPy.
582
583 Returns
584 -------
585 NumericArray
586 Rounded values of the NumericArray.
587
588 See Also
589 --------
590 numpy.around : Round values of an np.array.
591 DataFrame.round : Round values of a DataFrame.
592 Series.round : Round values of a Series.
593 """
594 if self.dtype.kind == "b":
595 return self
596 nv.validate_round(args, kwargs)
597 values = np.round(self._data, decimals=decimals, **kwargs)
598
599 # Usually we'll get same type as self, but ndarray[bool] casts to float
600 return self._maybe_mask_result(values, self._mask.copy())
601
602 # ------------------------------------------------------------------
603 # Unary Methods
604
605 def __invert__(self) -> Self:
606 return self._simple_new(~self._data, self._mask.copy())
607
608 def __neg__(self) -> Self:
609 return self._simple_new(-self._data, self._mask.copy())
610
611 def __pos__(self) -> Self:
612 return self.copy()
613
614 def __abs__(self) -> Self:
615 return self._simple_new(abs(self._data), self._mask.copy())
616
617 # ------------------------------------------------------------------
618
619 def _values_for_json(self) -> np.ndarray:
620 return np.asarray(self, dtype=object)
621
622 def to_numpy(
623 self,
624 dtype: npt.DTypeLike | None = None,
625 copy: bool = False,
626 na_value: object = lib.no_default,
627 ) -> np.ndarray:
628 """
629 Convert to a NumPy Array.
630
631 By default converts to an object-dtype NumPy array. Specify the `dtype` and
632 `na_value` keywords to customize the conversion.
633
634 Parameters
635 ----------
636 dtype : dtype, default object
637 The numpy dtype to convert to.
638 copy : bool, default False
639 Whether to ensure that the returned value is a not a view on
640 the array. Note that ``copy=False`` does not *ensure* that
641 ``to_numpy()`` is no-copy. Rather, ``copy=True`` ensure that
642 a copy is made, even if not strictly necessary. This is typically
643 only possible when no missing values are present and `dtype`
644 is the equivalent numpy dtype.
645 na_value : scalar, optional
646 Scalar missing value indicator to use in numpy array. Defaults
647 to the native missing value indicator of this array (pd.NA).
648
649 Returns
650 -------
651 numpy.ndarray
652
653 Examples
654 --------
655 An object-dtype is the default result
656
657 >>> a = pd.array([True, False, pd.NA], dtype="boolean")
658 >>> a.to_numpy()
659 array([True, False, <NA>], dtype=object)
660
661 When no missing values are present, an equivalent dtype can be used.
662
663 >>> pd.array([True, False], dtype="boolean").to_numpy(dtype="bool")
664 array([ True, False])
665 >>> pd.array([1, 2], dtype="Int64").to_numpy("int64")
666 array([1, 2])
667
668 However, requesting such dtype will raise a ValueError if
669 missing values are present and the default missing value :attr:`NA`
670 is used.
671
672 >>> a = pd.array([True, False, pd.NA], dtype="boolean")
673 >>> a
674 <BooleanArray>
675 [True, False, <NA>]
676 Length: 3, dtype: boolean
677
678 >>> a.to_numpy(dtype="bool")
679 Traceback (most recent call last):
680 ...
681 ValueError: cannot convert to bool numpy array in presence of missing values
682
683 Specify a valid `na_value` instead
684
685 >>> a.to_numpy(dtype="bool", na_value=False)
686 array([ True, False, False])
687 """
688 hasna = self._hasna
689 dtype, na_value = to_numpy_dtype_inference(self, dtype, na_value, hasna)
690 if dtype is None:
691 dtype = np.dtype(object)
692
693 if hasna:
694 if (
695 dtype != np.dtype(object)
696 and not is_string_dtype(dtype)
697 and na_value is libmissing.NA
698 ):
699 raise ValueError(
700 f"cannot convert to '{dtype}'-dtype NumPy array "
701 "with missing values. Specify an appropriate 'na_value' "
702 "for this dtype."
703 )
704 # don't pass copy to astype -> always need a copy since we are mutating
705 with warnings.catch_warnings():
706 warnings.filterwarnings("ignore", category=RuntimeWarning)
707 data = self._data.astype(dtype)
708 data[self._mask] = na_value
709 else:
710 with warnings.catch_warnings():
711 warnings.filterwarnings("ignore", category=RuntimeWarning)
712 data = self._data.astype(dtype, copy=copy)
713 if self._readonly and not copy and astype_is_view(self.dtype, dtype):
714 data = data.view()
715 data.flags.writeable = False
716 return data
717
718 def tolist(self) -> list[Any]:
719 """
720 Return a list of the values.
721
722 These are each a scalar type, which is a Python scalar
723 (for str, int, float) or a pandas scalar
724 (for Timestamp/Timedelta/Interval/Period)
725
726 Returns
727 -------
728 list
729 Python list of values in array.
730
731 See Also
732 --------
733 Index.to_list: Return a list of the values in the Index.
734 Series.to_list: Return a list of the values in the Series.
735
736 Examples
737 --------
738 >>> arr = pd.array([1, 2, 3])
739 >>> arr.tolist()
740 [1, 2, 3]
741 """
742 if self.ndim > 1:
743 return [x.tolist() for x in self]
744 dtype = None if self._hasna else self._data.dtype
745 return self.to_numpy(dtype=dtype, na_value=libmissing.NA).tolist()
746
747 @overload
748 def astype(self, dtype: npt.DTypeLike, copy: bool = ...) -> np.ndarray: ...
749
750 @overload
751 def astype(self, dtype: ExtensionDtype, copy: bool = ...) -> ExtensionArray: ...
752
753 @overload
754 def astype(self, dtype: AstypeArg, copy: bool = ...) -> ArrayLike: ...
755
756 def astype(self, dtype: AstypeArg, copy: bool = True) -> ArrayLike:
757 dtype = pandas_dtype(dtype)
758
759 if dtype == self.dtype:
760 if copy:
761 return self.copy()
762 return self
763
764 # if we are astyping to another nullable masked dtype, we can fastpath
765 if isinstance(dtype, BaseMaskedDtype):
766 # TODO deal with NaNs for FloatingArray case
767 with warnings.catch_warnings():
768 warnings.filterwarnings("ignore", category=RuntimeWarning)
769 # TODO: Is rounding what we want long term?
770 data = self._data.astype(dtype.numpy_dtype, copy=copy)
771 # mask is copied depending on whether the data was copied, and
772 # not directly depending on the `copy` keyword
773 mask = self._mask if data is self._data else self._mask.copy()
774 cls = dtype.construct_array_type()
775 return cls(data, mask, copy=False)
776
777 if isinstance(dtype, ExtensionDtype):
778 eacls = dtype.construct_array_type()
779 return eacls._from_sequence(self, dtype=dtype, copy=copy)
780
781 na_value: float | np.datetime64 | lib.NoDefault
782
783 # coerce
784 if dtype.kind == "f":
785 # In astype, we consider dtype=float to also mean na_value=np.nan
786 na_value = np.nan
787 elif dtype.kind == "M":
788 unit = np.datetime_data(dtype)[0]
789 na_value = np.datetime64("NaT", unit) # type: ignore[call-overload]
790 else:
791 na_value = lib.no_default
792
793 # to_numpy will also raise, but we get somewhat nicer exception messages here
794 if dtype.kind in "iu" and self._hasna:
795 raise ValueError("cannot convert NA to integer")
796 if dtype.kind == "b" and self._hasna:
797 # careful: astype_nansafe converts np.nan to True
798 raise ValueError("cannot convert float NaN to bool")
799
800 data = self.to_numpy(dtype=dtype, na_value=na_value, copy=copy)
801 return data
802
803 __array_priority__ = 1000 # higher than ndarray so ops dispatch to us
804
805 def __array__(
806 self, dtype: NpDtype | None = None, copy: bool | None = None
807 ) -> np.ndarray:
808 """
809 the array interface, return my values
810 We return an object array here to preserve our scalar values
811 """
812 if copy is False:
813 if not self._hasna:
814 # special case, here we can simply return the underlying data
815 result = np.array(self._data, dtype=dtype, copy=copy)
816 # If the ExtensionArray is readonly, make the numpy array readonly too
817 if self._readonly:
818 result = result.view()
819 result.flags.writeable = False
820 return result
821 raise ValueError(
822 "Unable to avoid copy while creating an array as requested."
823 )
824
825 if copy is None:
826 copy = False # The NumPy copy=False meaning is different here.
827 return self.to_numpy(dtype=dtype, copy=copy)
828
829 _HANDLED_TYPES: tuple[type, ...]
830
831 def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
832 # For MaskedArray inputs, we apply the ufunc to ._data
833 # and mask the result.
834
835 # this path never reaches ExtensionArray.__array_ufunc__, and a datetimelike
836 # scalar is not in _HANDLED_TYPES, so this has to precede that loop
837 ops.disallow_datetimelike_logical_ufunc(ufunc, inputs)
838
839 out = kwargs.get("out", ())
840
841 for x in inputs + out:
842 if not isinstance(x, (*self._HANDLED_TYPES, BaseMaskedArray)):
843 return NotImplemented
844
845 # for binary ops, use our custom dunder methods
846 result = arraylike.maybe_dispatch_ufunc_to_dunder_op(
847 self, ufunc, method, *inputs, **kwargs
848 )
849 if result is not NotImplemented:
850 return result
851
852 if "out" in kwargs:
853 # e.g. test_ufunc_with_out
854 return arraylike.dispatch_ufunc_with_out(
855 self, ufunc, method, *inputs, **kwargs
856 )
857
858 if method == "reduce":
859 result = arraylike.dispatch_reduction_ufunc(
860 self, ufunc, method, *inputs, **kwargs
861 )
862 if result is not NotImplemented:
863 return result
864
865 mask = np.zeros(len(self), dtype=bool)
866 inputs2 = []
867 for x in inputs:
868 if isinstance(x, BaseMaskedArray):
869 mask |= x._mask
870 inputs2.append(x._data)
871 else:
872 inputs2.append(x)
873
874 def reconstruct(x: np.ndarray):
875 # we don't worry about scalar `x` here, since we
876 # raise for reduce up above.
877 from pandas.core.arrays import (
878 BooleanArray,
879 FloatingArray,
880 IntegerArray,
881 )
882
883 if x.dtype.kind == "b":
884 m = mask.copy()
885 return BooleanArray(x, m)
886 elif x.dtype.kind in "iu":
887 m = mask.copy()
888 return IntegerArray(x, m)
889 elif x.dtype.kind == "f":
890 m = mask.copy()
891 if x.dtype == np.float16:
892 # reached in e.g. np.sqrt on BooleanArray
893 # we don't support float16
894 x = x.astype(np.float32)
895 if is_nan_na():
896 m[np.isnan(x)] = True
897 return FloatingArray(x, m)
898 else:
899 x[mask] = np.nan
900 return x
901
902 result = getattr(ufunc, method)(*inputs2, **kwargs)
903 if ufunc.nout > 1:
904 # e.g. np.divmod
905 return tuple(reconstruct(x) for x in result)
906 elif method == "reduce":
907 # e.g. np.add.reduce; test_ufunc_reduce_raises
908 if self._mask.any():
909 return self._na_value
910 return result
911 else:
912 return reconstruct(result)
913
914 def __arrow_array__(self, type=None):
915 """
916 Convert myself into a pyarrow Array.
917 """
918 import pyarrow as pa
919
920 return pa.array(self._data, mask=self._mask, type=type)
921
922 @property
923 def _hasna(self) -> bool:
924 # Note: this is expensive right now! The hope is that we can
925 # make this faster by having an optional mask, but not have to change
926 # source code using it..
927
928 return bool(self._mask.any())
929
930 def _propagate_mask(
931 self, mask: npt.NDArray[np.bool_] | None, other
932 ) -> npt.NDArray[np.bool_]:
933 if mask is None:
934 mask = self._mask.copy() # TODO: need test for BooleanArray needing a copy
935 if other is libmissing.NA:
936 # GH#45421 don't alter inplace
937 mask = mask | True
938 elif is_list_like(other) and len(other) == len(mask):
939 mask = mask | isna(other)
940 else:
941 mask = self._mask | mask
942 return mask
943
944 def _arith_method(self, other, op):
945 op_name = op.__name__
946 omask = None
947
948 if (
949 is_list_like(other)
950 and not isinstance(other, (list, np.ndarray, ExtensionArray))
951 and not ops.has_castable_attr(other)
952 ):
953 warnings.warn(
954 f"Operation with {type(other).__name__} is deprecated. "
955 "In a future version these will be treated as scalar-like. "
956 "To retain the old behavior, explicitly wrap in a Series "
957 "instead.",
958 Pandas4Warning,
959 stacklevel=find_stack_level(),
960 )
961
962 if (
963 not hasattr(other, "dtype")
964 and is_list_like(other)
965 and len(other) == len(self)
966 ):
967 # Try inferring masked dtype instead of casting to object
968 other = pd_array(other)
969 other = extract_array(other, extract_numpy=True)
970
971 if isinstance(other, BaseMaskedArray):
972 other, omask = other._data, other._mask
973
974 elif is_list_like(other):
975 if not isinstance(other, ExtensionArray):
976 other = np.asarray(other)
977 if other.ndim > 1:
978 raise NotImplementedError("can only perform ops with 1-d structures")
979
980 # We wrap the non-masked arithmetic logic used for numpy dtypes
981 # in Series/Index arithmetic ops.
982 other = ops.maybe_prepare_scalar_for_op(other, (len(self),))
983 pd_op = ops.get_array_op(op)
984 other = ensure_wrapped_if_datetimelike(other)
985
986 if isinstance(other, ExtensionArray) and isinstance(other.dtype, ArrowDtype):
987 # GH#58602
988 return NotImplemented
989
990 if op_name in {"pow", "rpow"} and isinstance(other, np.bool_):
991 # Avoid DeprecationWarning: In future, it will be an error
992 # for 'np.bool_' scalars to be interpreted as an index
993 # e.g. test_array_scalar_like_equivalence
994 other = bool(other)
995
996 mask = self._propagate_mask(omask, other)
997
998 if other is libmissing.NA:
999 result = np.ones_like(self._data)
1000 if self.dtype.kind == "b":
1001 if op_name in {
1002 "floordiv",
1003 "rfloordiv",
1004 "pow",
1005 "rpow",
1006 "truediv",
1007 "rtruediv",
1008 }:
1009 # GH#41165 Try to match non-masked Series behavior
1010 # This is still imperfect GH#46043
1011 raise NotImplementedError(
1012 f"operator '{op_name}' not implemented for bool dtypes"
1013 )
1014 if op_name in {"mod", "rmod"}:
1015 dtype = "int8"
1016 else:
1017 dtype = "bool"
1018 result = result.astype(dtype)
1019 elif "truediv" in op_name and self.dtype.kind != "f":
1020 # The actual data here doesn't matter since the mask
1021 # will be all-True, but since this is division, we want
1022 # to end up with floating dtype.
1023 result = result.astype(np.float64)
1024 elif op_name in {"divmod", "rdivmod"}:
1025 # GH#62196
1026 res = self._maybe_mask_result(result, mask)
1027 return res, res.copy()
1028 else:
1029 # Make sure we do this before the "pow" mask checks
1030 # to get an expected exception message on shape mismatch.
1031 if self.dtype.kind in "iu" and op_name in ["floordiv", "mod"]:
1032 # TODO(GH#30188) ATM we don't match the behavior of non-masked
1033 # types with respect to floordiv-by-zero
1034 pd_op = op
1035
1036 with np.errstate(all="ignore"):
1037 result = pd_op(self._data, other)
1038
1039 if op_name == "pow":
1040 # 1 ** x is 1.
1041 mask = np.where((self._data == 1) & ~self._mask, False, mask)
1042 # x ** 0 is 1.
1043 if omask is not None:
1044 mask = np.where((other == 0) & ~omask, False, mask)
1045 elif other is not libmissing.NA:
1046 mask = np.where(other == 0, False, mask)
1047
1048 elif op_name == "rpow":
1049 # 1 ** x is 1.
1050 if omask is not None:
1051 mask = np.where((other == 1) & ~omask, False, mask)
1052 elif other is not libmissing.NA:
1053 mask = np.where(other == 1, False, mask)
1054 # x ** 0 is 1.
1055 mask = np.where((self._data == 0) & ~self._mask, False, mask)
1056
1057 return self._maybe_mask_result(result, mask)
1058
1059 _logical_method = _arith_method
1060
1061 def _cmp_method(self, other, op) -> BooleanArray:
1062 from pandas.core.arrays import BooleanArray
1063
1064 mask = None
1065
1066 if isinstance(other, ExtensionArray) and isinstance(other.dtype, ArrowDtype):
1067 # GH#58602
1068 return NotImplemented
1069
1070 elif isinstance(other, BaseMaskedArray):
1071 other, mask = other._data, other._mask
1072
1073 elif is_list_like(other):
1074 if not isinstance(
1075 other, (list, np.ndarray, ExtensionArray)
1076 ) and not ops.has_castable_attr(other):
1077 warnings.warn(
1078 f"Operation with {type(other).__name__} is deprecated. "
1079 "In a future version these will be treated as scalar-like. "
1080 "To retain the old behavior, explicitly wrap in a Series "
1081 "instead.",
1082 Pandas4Warning,
1083 stacklevel=find_stack_level(),
1084 )
1085 other = np.asarray(other)
1086 if other.ndim > 1:
1087 raise NotImplementedError("can only perform ops with 1-d structures")
1088 if len(self) != len(other):
1089 raise ValueError("Lengths must match to compare")
1090
1091 if other is libmissing.NA:
1092 # numpy does not handle pd.NA well as "other" scalar (it returns
1093 # a scalar False instead of an array)
1094 # This may be fixed by NA.__array_ufunc__. Revisit this check
1095 # once that's implemented.
1096 result = np.zeros(self._data.shape, dtype="bool")
1097 mask = np.ones(self._data.shape, dtype="bool")
1098 else:
1099 with warnings.catch_warnings():
1100 # numpy may show a FutureWarning or DeprecationWarning:
1101 # elementwise comparison failed; returning scalar instead,
1102 # but in the future will perform elementwise comparison
1103 # before returning NotImplemented. We fall back to the correct
1104 # behavior today, so that should be fine to ignore.
1105 warnings.filterwarnings("ignore", "elementwise", FutureWarning)
1106 warnings.filterwarnings("ignore", "elementwise", DeprecationWarning)
1107 method = getattr(self._data, f"__{op.__name__}__")
1108 result = method(other)
1109
1110 if result is NotImplemented:
1111 result = invalid_comparison(self._data, other, op)
1112
1113 mask = self._propagate_mask(mask, other)
1114 return BooleanArray(result, mask, copy=False)
1115
1116 def _maybe_mask_result(
1117 self, result: np.ndarray | tuple[np.ndarray, np.ndarray], mask: np.ndarray
1118 ):
1119 """
1120 Parameters
1121 ----------
1122 result : array-like or tuple[array-like]
1123 mask : array-like bool
1124 """
1125 if isinstance(result, tuple):
1126 # i.e. divmod
1127 div, mod = result
1128 return (
1129 self._maybe_mask_result(div, mask),
1130 self._maybe_mask_result(mod, mask),
1131 )
1132
1133 if result.dtype.kind == "f":
1134 from pandas.core.arrays import FloatingArray
1135
1136 if is_nan_na():
1137 mask[np.isnan(result)] = True
1138
1139 return FloatingArray(result, mask, copy=False)
1140
1141 elif result.dtype.kind == "b":
1142 from pandas.core.arrays import BooleanArray
1143
1144 return BooleanArray(result, mask, copy=False)
1145
1146 elif lib.is_np_dtype(result.dtype, "m") and is_supported_dtype(result.dtype):
1147 # e.g. test_numeric_arr_mul_tdscalar_numexpr_path
1148 from pandas.core.arrays import TimedeltaArray
1149
1150 unit = np.datetime_data(result.dtype)[0]
1151 result[mask] = np.timedelta64("NaT", unit) # type: ignore[call-overload]
1152
1153 if not isinstance(result, TimedeltaArray):
1154 return TimedeltaArray._simple_new(result, dtype=result.dtype)
1155
1156 return result
1157
1158 elif result.dtype.kind in "iu":
1159 from pandas.core.arrays import IntegerArray
1160
1161 return IntegerArray(result, mask, copy=False)
1162
1163 elif result.dtype == object:
1164 result[mask] = self.dtype.na_value
1165 return result
1166 else:
1167 result[mask] = np.nan
1168 return result
1169
1170 def isna(self) -> np.ndarray:
1171 return self._mask.copy()
1172
1173 @property
1174 def _na_value(self):
1175 return self.dtype.na_value
1176
1177 @property
1178 def nbytes(self) -> int:
1179 return self._data.nbytes + self._mask.nbytes
1180
1181 @classmethod
1182 def _concat_same_type(
1183 cls,
1184 to_concat: Sequence[Self],
1185 axis: AxisInt = 0,
1186 ) -> Self:
1187 data = np.concatenate([x._data for x in to_concat], axis=axis)
1188 mask = np.concatenate([x._mask for x in to_concat], axis=axis)
1189 return cls(data, mask)
1190
1191 def _hash_pandas_object(
1192 self, *, encoding: str, hash_key: str, categorize: bool
1193 ) -> npt.NDArray[np.uint64]:
1194 hashed_array = hash_array(
1195 self._data, encoding=encoding, hash_key=hash_key, categorize=categorize
1196 )
1197 hashed_array[self.isna()] = hash(self.dtype.na_value)
1198 return hashed_array
1199
1200 def take(
1201 self,
1202 indexer,
1203 *,
1204 allow_fill: bool = False,
1205 fill_value: Scalar | None = None,
1206 axis: AxisInt = 0,
1207 ) -> Self:
1208 # we always fill with 1 internally
1209 # to avoid upcasting
1210 data_fill_value = (
1211 self.dtype._internal_fill_value if isna(fill_value) else fill_value
1212 )
1213 result = take(
1214 self._data,
1215 indexer,
1216 fill_value=data_fill_value,
1217 allow_fill=allow_fill,
1218 axis=axis,
1219 )
1220
1221 mask = take(
1222 self._mask, indexer, fill_value=True, allow_fill=allow_fill, axis=axis
1223 )
1224
1225 # if we are filling
1226 # we only fill where the indexer is null
1227 # not existing missing values
1228 # TODO(jreback) what if we have a non-na float as a fill value?
1229 if allow_fill and notna(fill_value):
1230 fill_mask = np.asarray(indexer) == -1
1231 result[fill_mask] = fill_value
1232 mask = mask ^ fill_mask
1233
1234 return self._simple_new(result, mask)
1235
1236 # error: Return type "BooleanArray" of "isin" incompatible with return type
1237 # "ndarray" in supertype "ExtensionArray"
1238 def isin(self, values: ArrayLike) -> BooleanArray: # type: ignore[override]
1239 from pandas.core.arrays import BooleanArray
1240
1241 # algorithms.isin will eventually convert values to an ndarray, so no extra
1242 # cost to doing it here first
1243 values_arr = np.asarray(values)
1244 result = isin(self._data, values_arr)
1245
1246 if self._hasna:
1247 values_have_NA = values_arr.dtype == object and any(
1248 val is self.dtype.na_value for val in values_arr
1249 )
1250
1251 # For now, NA does not propagate so set result according to presence of NA,
1252 # see https://github.com/pandas-dev/pandas/pull/38379 for some discussion
1253 result[self._mask] = values_have_NA
1254
1255 mask = np.zeros(self._data.shape, dtype=bool)
1256 return BooleanArray(result, mask, copy=False)
1257
1258 def copy(self) -> Self:
1259 data = self._data.copy()
1260 mask = self._mask.copy()
1261 return self._simple_new(data, mask)
1262
1263 @overload
1264 def view(self, dtype: None = ...) -> Self: ...
1265
1266 @overload
1267 def view(self, dtype: Dtype | None = ...) -> ArrayLike: ...
1268
1269 def view(self, dtype: Dtype | None = None) -> ArrayLike:
1270 if dtype is not None:
1271 return super().view(dtype)
1272 result = self._simple_new(self._data[:], self._mask[:])
1273 result._readonly = self._readonly
1274 return result
1275
1276 @property
1277 def _is_monotonic_increasing(self) -> bool:
1278 if self._hasna:
1279 return False
1280 data = self._data
1281 if data.dtype.kind == "b":
1282 data = data.view("uint8")
1283 return libalgos.is_monotonic(data, timelike=False)[0]
1284
1285 @property
1286 def _is_monotonic_decreasing(self) -> bool:
1287 if self._hasna:
1288 return False
1289 data = self._data
1290 if data.dtype.kind == "b":
1291 data = data.view("uint8")
1292 return libalgos.is_monotonic(data, timelike=False)[1]
1293
1294 def _rank(
1295 self,
1296 *,
1297 axis: AxisInt = 0,
1298 method: RankMethod = "average",
1299 na_option: RankNaOption = "keep",
1300 ascending: bool = True,
1301 pct: bool = False,
1302 ):
1303 # GH#62043 Avoid going through copy-making ensure_data in algorithms.rank
1304 if axis != 0 or self.ndim != 1:
1305 raise NotImplementedError
1306
1307 from pandas.core.arrays import FloatingArray
1308
1309 data = self._data
1310 if data.dtype.kind == "b":
1311 data = data.view("uint8")
1312
1313 result = libalgos.rank_1d(
1314 data,
1315 is_datetimelike=False,
1316 ties_method=method,
1317 ascending=ascending,
1318 na_option=na_option,
1319 pct=pct,
1320 mask=self.isna(),
1321 )
1322 if na_option in ["top", "bottom"]:
1323 mask = np.zeros(self.shape, dtype=bool)
1324 else:
1325 mask = self._mask.copy()
1326
1327 if method != "average" and not pct:
1328 if na_option not in ["top", "bottom"]:
1329 result[self._mask] = 0 # avoid warning on casting
1330 result = result.astype("uint64", copy=False)
1331 from pandas.core.arrays import IntegerArray
1332
1333 return IntegerArray(result, mask=mask)
1334
1335 return FloatingArray(result, mask=mask)
1336
1337 def duplicated(
1338 self, keep: Literal["first", "last", False] = "first"
1339 ) -> npt.NDArray[np.bool_]:
1340 """
1341 Return boolean ndarray denoting duplicate values.
1342
1343 Parameters
1344 ----------
1345 keep : {'first', 'last', False}, default 'first'
1346 - ``first`` : Mark duplicates as ``True`` except for the first occurrence.
1347 - ``last`` : Mark duplicates as ``True`` except for the last occurrence.
1348 - False : Mark all duplicates as ``True``.
1349
1350 Returns
1351 -------
1352 ndarray[bool]
1353 With true in indices where elements are duplicated and false otherwise.
1354
1355 See Also
1356 --------
1357 DataFrame.duplicated : Return boolean Series denoting
1358 duplicate rows.
1359 Series.duplicated : Indicate duplicate Series values.
1360 api.extensions.ExtensionArray.unique : Compute the ExtensionArray
1361 of unique values.
1362
1363 Examples
1364 --------
1365 >>> pd.array([1, 1, 2, 3, 3], dtype="Int64").duplicated()
1366 array([False, True, False, False, True])
1367 """
1368 values = self._data
1369 mask = self._mask
1370 return algos.duplicated(values, keep=keep, mask=mask)
1371
1372 def unique(self) -> Self:
1373 """
1374 Compute the BaseMaskedArray of unique values.
1375
1376 Returns
1377 -------
1378 uniques : BaseMaskedArray
1379 """
1380 uniques, mask = algos.unique_with_mask(self._data, self._mask)
1381 return self._simple_new(uniques, mask)
1382
1383 def searchsorted(
1384 self,
1385 value: NumpyValueArrayLike | ExtensionArray,
1386 side: Literal["left", "right"] = "left",
1387 sorter: NumpySorter | None = None,
1388 ) -> npt.NDArray[np.intp] | np.intp:
1389 """
1390 Find indices where elements should be inserted to maintain order.
1391
1392 Find the indices into a sorted array `self` (a) such that, if the
1393 corresponding elements in `value` were inserted before the indices,
1394 the order of `self` would be preserved.
1395
1396 Assuming that `self` is sorted:
1397
1398 ====== ================================
1399 `side` returned index `i` satisfies
1400 ====== ================================
1401 left ``self[i-1] < value <= self[i]``
1402 right ``self[i-1] <= value < self[i]``
1403 ====== ================================
1404
1405 Parameters
1406 ----------
1407 value : array-like, list or scalar
1408 Value(s) to insert into `self`.
1409 side : {'left', 'right'}, optional
1410 If 'left', the index of the first suitable location found is given.
1411 If 'right', return the last such index. If there is no suitable
1412 index, return either 0 or N (where N is the length of `self`).
1413 sorter : 1-D array-like, optional
1414 Optional array of integer indices that sort array a into ascending
1415 order. They are typically the result of argsort.
1416
1417 Returns
1418 -------
1419 array of ints or int
1420 If value is array-like, array of insertion points.
1421 If value is scalar, a single integer.
1422
1423 See Also
1424 --------
1425 numpy.searchsorted : Similar method from NumPy.
1426
1427 Examples
1428 --------
1429 >>> arr = pd.array([1, 2, 3, 5])
1430 >>> arr.searchsorted([4])
1431 array([3])
1432 """
1433 if self._hasna:
1434 raise ValueError(
1435 "searchsorted requires array to be sorted, which is impossible "
1436 "with NAs present."
1437 )
1438 if isinstance(value, ExtensionArray):
1439 value = value.astype(object)
1440 # Base class searchsorted would cast to object, which is *much* slower.
1441 return self._data.searchsorted(value, side=side, sorter=sorter)
1442
1443 def factorize(
1444 self,
1445 use_na_sentinel: bool = True,
1446 ) -> tuple[np.ndarray, ExtensionArray]:
1447 """
1448 Encode the extension array as an enumerated type.
1449
1450 Parameters
1451 ----------
1452 use_na_sentinel : bool, default True
1453 If True, the sentinel -1 will be used for NaN values. If False,
1454 NaN values will be encoded as non-negative integers and will not drop the
1455 NaN from the uniques of the values.
1456
1457 Returns
1458 -------
1459 codes : ndarray
1460 An integer NumPy array that's an indexer into the original
1461 ExtensionArray.
1462 uniques : ExtensionArray
1463 An ExtensionArray containing the unique values of `self`.
1464
1465 .. note::
1466
1467 uniques will *not* contain an entry for the NA value of
1468 the ExtensionArray if there are any missing values present
1469 in `self`.
1470
1471 See Also
1472 --------
1473 factorize : Top-level factorize method that dispatches here.
1474
1475 Notes
1476 -----
1477 :meth:`pandas.factorize` offers a `sort` keyword as well.
1478
1479 Examples
1480 --------
1481 >>> idx1 = pd.PeriodIndex(
1482 ... ["2014-01", "2014-01", "2014-02", "2014-02", "2014-03", "2014-03"],
1483 ... freq="M",
1484 ... )
1485 >>> arr, idx = idx1.factorize()
1486 >>> arr
1487 array([0, 0, 1, 1, 2, 2])
1488 >>> idx
1489 PeriodIndex(['2014-01', '2014-02', '2014-03'], dtype='period[M]')
1490 """
1491 arr = self._data
1492 mask = self._mask
1493
1494 # Use a sentinel for na; recode and add NA to uniques if necessary below
1495 codes, uniques = factorize_array(arr, use_na_sentinel=True, mask=mask)
1496
1497 # check that factorize_array correctly preserves dtype.
1498 assert uniques.dtype == self.dtype.numpy_dtype, (uniques.dtype, self.dtype)
1499
1500 has_na = mask.any()
1501 if use_na_sentinel or not has_na:
1502 size = len(uniques)
1503 else:
1504 # Make room for an NA value
1505 size = len(uniques) + 1
1506 uniques_mask = np.zeros(size, dtype=bool)
1507 if not use_na_sentinel and has_na:
1508 na_index = mask.argmax()
1509 # Insert na with the proper code
1510 if na_index == 0:
1511 na_code = np.intp(0)
1512 else:
1513 na_code = codes[:na_index].max() + 1
1514 codes[codes >= na_code] += 1
1515 codes[codes == -1] = na_code
1516 # dummy value for uniques; not used since uniques_mask will be True
1517 uniques = np.insert(uniques, na_code, 0)
1518 uniques_mask[na_code] = True
1519 uniques_ea = self._simple_new(uniques, uniques_mask)
1520
1521 return codes, uniques_ea
1522
1523 def _values_for_argsort(self) -> np.ndarray:
1524 """
1525 Return values for sorting.
1526
1527 Returns
1528 -------
1529 ndarray
1530 The transformed values should maintain the ordering between values
1531 within the array.
1532
1533 See Also
1534 --------
1535 ExtensionArray.argsort : Return the indices that would sort this array.
1536
1537 Notes
1538 -----
1539 The caller is responsible for *not* modifying these values in-place, so
1540 it is safe for implementers to give views on ``self``.
1541
1542 Functions that use this (e.g. ``ExtensionArray.argsort``) should ignore
1543 entries with missing values in the original array (according to
1544 ``self.isna()``). This means that the corresponding entries in the returned
1545 array don't need to be modified to sort correctly.
1546
1547 Examples
1548 --------
1549 In most cases, this is the underlying Numpy array of the ``ExtensionArray``:
1550
1551 >>> arr = pd.array([1, 2, 3])
1552 >>> arr._values_for_argsort()
1553 array([1, 2, 3])
1554 """
1555 return self._data
1556
1557 def value_counts(self, dropna: bool = True) -> Series:
1558 """
1559 Returns a Series containing counts of each unique value.
1560
1561 Parameters
1562 ----------
1563 dropna : bool, default True
1564 Don't include counts of missing values.
1565
1566 Returns
1567 -------
1568 counts : Series
1569
1570 See Also
1571 --------
1572 Series.value_counts
1573 """
1574 from pandas import (
1575 Index,
1576 Series,
1577 )
1578 from pandas.arrays import IntegerArray
1579
1580 keys, value_counts, na_counter = algos.value_counts_arraylike(
1581 self._data, dropna=dropna, mask=self._mask
1582 )
1583 mask_index = np.zeros((len(value_counts),), dtype=np.bool_)
1584 mask = mask_index.copy()
1585
1586 if na_counter > 0:
1587 mask_index[-1] = True
1588
1589 arr = IntegerArray(value_counts, mask)
1590 index = Index(
1591 self.dtype.construct_array_type()(
1592 keys, # type: ignore[arg-type]
1593 mask_index,
1594 ),
1595 copy=False,
1596 )
1597 return Series(arr, index=index, name="count", copy=False)
1598
1599 def _mode(self, dropna: bool = True) -> Self:
1600 result, res_mask = mode(self._data, dropna=dropna, mask=self._mask)
1601 result = type(self)(result, res_mask)
1602 return result[result.argsort()]
1603
1604 def equals(self, other) -> bool:
1605 """
1606 Return if another array is equivalent to this array.
1607
1608 Equivalent means that both arrays have the same shape and dtype, and
1609 all values compare equal. Missing values in the same location are
1610 considered equal (in contrast with normal equality).
1611
1612 Parameters
1613 ----------
1614 other : ExtensionArray
1615 Array to compare to this Array.
1616
1617 Returns
1618 -------
1619 boolean
1620 Whether the arrays are equivalent.
1621
1622 See Also
1623 --------
1624 numpy.array_equal : Equivalent method for numpy array.
1625 Series.equals : Equivalent method for Series.
1626 DataFrame.equals : Equivalent method for DataFrame.
1627
1628 Examples
1629 --------
1630 >>> arr1 = pd.array([1, 2, np.nan])
1631 >>> arr2 = pd.array([1, 2, np.nan])
1632 >>> arr1.equals(arr2)
1633 True
1634
1635 >>> arr1 = pd.array([1, 3, np.nan])
1636 >>> arr2 = pd.array([1, 2, np.nan])
1637 >>> arr1.equals(arr2)
1638 False
1639 """
1640 if type(self) != type(other):
1641 return False
1642 if other.dtype != self.dtype:
1643 return False
1644
1645 # GH#44382 if e.g. self[1] is np.nan and other[1] is pd.NA, we are NOT
1646 # equal.
1647 if not lib.array_equivalent_bytes(self._mask, other._mask):
1648 return False
1649
1650 left = self._data[~self._mask]
1651 right = other._data[~other._mask]
1652 return array_equivalent(left, right, strict_nan=True, dtype_equal=True)
1653
1654 def _quantile(
1655 self, qs: npt.NDArray[np.float64], interpolation: str
1656 ) -> BaseMaskedArray:
1657 """
1658 Dispatch to quantile_with_mask, needed because we do not have
1659 _from_factorized.
1660
1661 Notes
1662 -----
1663 We assume that all impacted cases are 1D-only.
1664 """
1665 res = quantile_with_mask(
1666 self._data,
1667 mask=self._mask,
1668 # TODO(GH#40932): na_value_for_dtype(self.dtype.numpy_dtype)
1669 # instead of np.nan
1670 fill_value=np.nan,
1671 qs=qs,
1672 interpolation=interpolation,
1673 )
1674
1675 if self._hasna:
1676 # Our result mask is all-False unless we are all-NA, in which
1677 # case it is all-True.
1678 if self.ndim == 2:
1679 # I think this should be out_mask=self.isna().all(axis=1)
1680 # but am holding off until we have tests
1681 raise NotImplementedError
1682 if self.isna().all():
1683 out_mask = np.ones(res.shape, dtype=bool)
1684
1685 if is_integer_dtype(self.dtype):
1686 # We try to maintain int dtype if possible for not all-na case
1687 # as well
1688 res = np.zeros(res.shape, dtype=self.dtype.numpy_dtype)
1689 else:
1690 out_mask = np.zeros(res.shape, dtype=bool)
1691 else:
1692 out_mask = np.zeros(res.shape, dtype=bool)
1693 return self._maybe_mask_result(res, mask=out_mask)
1694
1695 # ------------------------------------------------------------------
1696 # Reductions
1697
1698 def _reduce(
1699 self,
1700 name: str,
1701 *,
1702 skipna: bool = True,
1703 keepdims: bool = False,
1704 axis: int | None = 0,
1705 **kwargs,
1706 ):
1707 result: Any
1708 if name == "count":
1709 result = self.count()
1710 elif name in {"any", "all", "min", "max", "sum", "prod", "mean", "var", "std"}:
1711 result = getattr(self, name)(skipna=skipna, axis=axis, **kwargs)
1712 else:
1713 # median, skew, kurt, sem
1714 data = self._data
1715 mask = self._mask
1716 op = getattr(nanops, f"nan{name}")
1717 result = op(data, axis=axis, skipna=skipna, mask=mask, **kwargs)
1718
1719 if self.ndim == 2 and axis is not None:
1720 if keepdims:
1721 raise NotImplementedError(
1722 "keepdims=True is not implemented for 2D MaskedArray "
1723 "reductions with axis"
1724 )
1725 return result
1726
1727 if keepdims:
1728 if isna(result):
1729 return self._wrap_na_result(name=name, axis=0, mask_size=(1,))
1730 else:
1731 if using_python_scalars():
1732 result = np.array([result])
1733 else:
1734 result = result.reshape(1)
1735 mask = np.zeros(1, dtype=bool)
1736 return self._maybe_mask_result(result, mask)
1737
1738 if isna(result):
1739 return libmissing.NA
1740 else:
1741 return result
1742
1743 def _wrap_reduction_result(self, name: str, result, *, skipna, axis):
1744 if isinstance(result, np.ndarray):
1745 if skipna:
1746 # we only retain mask for all-NA rows/columns
1747 mask = self._mask.all(axis=axis)
1748 else:
1749 mask = self._mask.any(axis=axis)
1750
1751 return self._maybe_mask_result(result, mask)
1752 return result
1753
1754 def _wrap_na_result(self, *, name, axis, mask_size):
1755 mask = np.ones(mask_size, dtype=bool)
1756
1757 float_dtyp = "float32" if self.dtype == "Float32" else "float64"
1758 if name in ("any", "all"):
1759 np_dtype = "bool"
1760 elif name in ["mean", "median", "var", "std", "skew", "kurt", "sem"]:
1761 np_dtype = float_dtyp
1762 elif name in ["min", "max"] or self.dtype.itemsize == 8:
1763 np_dtype = self.dtype.numpy_dtype.name
1764 else:
1765 is_windows_or_32bit = is_platform_windows() or not IS64
1766 int_dtyp = "int32" if is_windows_or_32bit else "int64"
1767 uint_dtyp = "uint32" if is_windows_or_32bit else "uint64"
1768 np_dtype = {"b": int_dtyp, "i": int_dtyp, "u": uint_dtyp, "f": float_dtyp}[
1769 self.dtype.kind
1770 ]
1771
1772 value = np.array([1], dtype=np_dtype)
1773 return self._maybe_mask_result(value, mask=mask)
1774
1775 def _wrap_min_count_reduction_result(
1776 self, name: str, result, *, skipna, min_count, axis
1777 ):
1778 if min_count == 0 and skipna and isinstance(result, np.ndarray):
1779 return self._maybe_mask_result(result, np.zeros(result.shape, dtype=bool))
1780 return self._wrap_reduction_result(name, result, skipna=skipna, axis=axis)
1781
1782 def sum(
1783 self,
1784 *,
1785 skipna: bool = True,
1786 min_count: int = 0,
1787 axis: AxisInt | None = 0,
1788 **kwargs,
1789 ):
1790 nv.validate_sum((), kwargs)
1791
1792 result = masked_reductions.sum(
1793 self._data,
1794 self._mask,
1795 skipna=skipna,
1796 min_count=min_count,
1797 axis=axis,
1798 )
1799 return self._wrap_min_count_reduction_result(
1800 "sum", result, skipna=skipna, min_count=min_count, axis=axis
1801 )
1802
1803 def prod(
1804 self,
1805 *,
1806 skipna: bool = True,
1807 min_count: int = 0,
1808 axis: AxisInt | None = 0,
1809 **kwargs,
1810 ):
1811 nv.validate_prod((), kwargs)
1812
1813 result = masked_reductions.prod(
1814 self._data,
1815 self._mask,
1816 skipna=skipna,
1817 min_count=min_count,
1818 axis=axis,
1819 )
1820 return self._wrap_min_count_reduction_result(
1821 "prod", result, skipna=skipna, min_count=min_count, axis=axis
1822 )
1823
1824 def mean(self, *, skipna: bool = True, axis: AxisInt | None = 0, **kwargs):
1825 nv.validate_mean((), kwargs)
1826 result = masked_reductions.mean(
1827 self._data,
1828 self._mask,
1829 skipna=skipna,
1830 axis=axis,
1831 )
1832 return self._wrap_reduction_result("mean", result, skipna=skipna, axis=axis)
1833
1834 def median(self, *, skipna: bool = True, axis: AxisInt | None = 0, **kwargs):
1835 nv.validate_median((), kwargs)
1836 result = self._reduce("median", skipna=skipna, axis=axis, **kwargs)
1837 return self._wrap_reduction_result("median", result, skipna=skipna, axis=axis)
1838
1839 def kurt(self, *, skipna: bool = True, axis: AxisInt | None = 0, **kwargs):
1840 nv.validate_stat_ddof_func((), kwargs, fname="kurt")
1841 result = self._reduce("kurt", skipna=skipna, axis=axis, **kwargs)
1842 return self._wrap_reduction_result("kurt", result, skipna=skipna, axis=axis)
1843
1844 def sem(
1845 self, *, skipna: bool = True, axis: AxisInt | None = 0, ddof: int = 1, **kwargs
1846 ):
1847 nv.validate_stat_ddof_func((), kwargs, fname="sem")
1848 result = self._reduce("sem", skipna=skipna, axis=axis, ddof=ddof, **kwargs)
1849 return self._wrap_reduction_result("sem", result, skipna=skipna, axis=axis)
1850
1851 def skew(self, *, skipna: bool = True, axis: AxisInt | None = 0, **kwargs):
1852 nv.validate_stat_ddof_func((), kwargs, fname="skew")
1853 result = self._reduce("skew", skipna=skipna, axis=axis, **kwargs)
1854 return self._wrap_reduction_result("skew", result, skipna=skipna, axis=axis)
1855
1856 def var(
1857 self, *, skipna: bool = True, axis: AxisInt | None = 0, ddof: int = 1, **kwargs
1858 ):
1859 nv.validate_stat_ddof_func((), kwargs, fname="var")
1860 result = masked_reductions.var(
1861 self._data,
1862 self._mask,
1863 skipna=skipna,
1864 axis=axis,
1865 ddof=ddof,
1866 )
1867 return self._wrap_reduction_result("var", result, skipna=skipna, axis=axis)
1868
1869 def std(
1870 self, *, skipna: bool = True, axis: AxisInt | None = 0, ddof: int = 1, **kwargs
1871 ):
1872 nv.validate_stat_ddof_func((), kwargs, fname="std")
1873 result = masked_reductions.std(
1874 self._data,
1875 self._mask,
1876 skipna=skipna,
1877 axis=axis,
1878 ddof=ddof,
1879 )
1880 return self._wrap_reduction_result("std", result, skipna=skipna, axis=axis)
1881
1882 def min(self, *, skipna: bool = True, axis: AxisInt | None = 0, **kwargs):
1883 nv.validate_min((), kwargs)
1884 result = masked_reductions.min(
1885 self._data,
1886 self._mask,
1887 skipna=skipna,
1888 axis=axis,
1889 )
1890 return self._wrap_reduction_result("min", result, skipna=skipna, axis=axis)
1891
1892 def max(self, *, skipna: bool = True, axis: AxisInt | None = 0, **kwargs):
1893 nv.validate_max((), kwargs)
1894 result = masked_reductions.max(
1895 self._data,
1896 self._mask,
1897 skipna=skipna,
1898 axis=axis,
1899 )
1900 return self._wrap_reduction_result("max", result, skipna=skipna, axis=axis)
1901
1902 def count(self) -> np.int64:
1903 return (~self._mask).sum()
1904
1905 @overload
1906 def any(
1907 self, *, skipna: Literal[True] = ..., axis: AxisInt | None = ..., **kwargs
1908 ) -> np.bool_: ...
1909
1910 @overload
1911 def any(
1912 self, *, skipna: bool, axis: AxisInt | None = ..., **kwargs
1913 ) -> np.bool_ | NAType: ...
1914
1915 def any(
1916 self, *, skipna: bool = True, axis: AxisInt | None = 0, **kwargs
1917 ) -> np.bool_ | NAType:
1918 """
1919 Return whether any element is truthy.
1920
1921 Returns False unless there is at least one element that is truthy.
1922 By default, NAs are skipped. If ``skipna=False`` is specified and
1923 missing values are present, similar :ref:`Kleene logic <boolean.kleene>`
1924 is used as for logical operations.
1925
1926 Parameters
1927 ----------
1928 skipna : bool, default True
1929 Exclude NA values. If the entire array is NA and `skipna` is
1930 True, then the result will be False, as for an empty array.
1931 If `skipna` is False, the result will still be True if there is
1932 at least one element that is truthy, otherwise NA will be returned
1933 if there are NA's present.
1934 axis : int, optional, default 0
1935 **kwargs : any, default None
1936 Additional keywords have no effect but might be accepted for
1937 compatibility with NumPy. See :ref:`gotchas.numpy_kwargs` for more.
1938
1939 Returns
1940 -------
1941 bool or :attr:`pandas.NA`
1942
1943 See Also
1944 --------
1945 numpy.any : Numpy version of this method.
1946 BaseMaskedArray.all : Return whether all elements are truthy.
1947
1948 Examples
1949 --------
1950 The result indicates whether any element is truthy (and by default
1951 skips NAs):
1952
1953 >>> pd.array([True, False, True]).any()
1954 np.True_
1955 >>> pd.array([True, False, pd.NA]).any()
1956 np.True_
1957 >>> pd.array([False, False, pd.NA]).any()
1958 np.False_
1959 >>> pd.array([], dtype="boolean").any()
1960 np.False_
1961 >>> pd.array([pd.NA], dtype="boolean").any()
1962 np.False_
1963 >>> pd.array([pd.NA], dtype="Float64").any()
1964 np.False_
1965
1966 With ``skipna=False``, the result can be NA if this is logically
1967 required (whether ``pd.NA`` is True or False influences the result):
1968
1969 >>> pd.array([True, False, pd.NA]).any(skipna=False)
1970 np.True_
1971 >>> pd.array([1, 0, pd.NA]).any(skipna=False)
1972 np.True_
1973 >>> pd.array([False, False, pd.NA]).any(skipna=False)
1974 <NA>
1975 >>> pd.array([0, 0, pd.NA]).any(skipna=False)
1976 <NA>
1977 """
1978 nv.validate_any((), kwargs)
1979
1980 values = self._data.copy()
1981 np.putmask(values, self._mask, self.dtype._falsey_value)
1982 result = values.any(axis=axis)
1983
1984 if isinstance(result, np.ndarray):
1985 if skipna:
1986 mask = np.zeros(result.shape, dtype=bool)
1987 else:
1988 # Kleene logic: False | NA = NA, True | NA = True
1989 mask = ~result & self._mask.any(axis=axis)
1990 return self._maybe_mask_result(result, mask)
1991
1992 if skipna:
1993 return result
1994 elif result or len(self) == 0 or not self._mask.any():
1995 return result
1996 else:
1997 return self.dtype.na_value
1998
1999 @overload
2000 def all(
2001 self, *, skipna: Literal[True] = ..., axis: AxisInt | None = ..., **kwargs
2002 ) -> np.bool_: ...
2003
2004 @overload
2005 def all(
2006 self, *, skipna: bool, axis: AxisInt | None = ..., **kwargs
2007 ) -> np.bool_ | NAType: ...
2008
2009 def all(
2010 self, *, skipna: bool = True, axis: AxisInt | None = 0, **kwargs
2011 ) -> np.bool_ | NAType:
2012 """
2013 Return whether all elements are truthy.
2014
2015 Returns True unless there is at least one element that is falsey.
2016 By default, NAs are skipped. If ``skipna=False`` is specified and
2017 missing values are present, similar :ref:`Kleene logic <boolean.kleene>`
2018 is used as for logical operations.
2019
2020 Parameters
2021 ----------
2022 skipna : bool, default True
2023 Exclude NA values. If the entire array is NA and `skipna` is
2024 True, then the result will be True, as for an empty array.
2025 If `skipna` is False, the result will still be False if there is
2026 at least one element that is falsey, otherwise NA will be returned
2027 if there are NA's present.
2028 axis : int, optional, default 0
2029 **kwargs : any, default None
2030 Additional keywords have no effect but might be accepted for
2031 compatibility with NumPy. See :ref:`gotchas.numpy_kwargs` for more.
2032
2033 Returns
2034 -------
2035 bool or :attr:`pandas.NA`
2036
2037 See Also
2038 --------
2039 numpy.all : Numpy version of this method.
2040 BooleanArray.any : Return whether any element is truthy.
2041
2042 Examples
2043 --------
2044 The result indicates whether all elements are truthy (and by default
2045 skips NAs):
2046
2047 >>> pd.array([True, True, pd.NA]).all()
2048 np.True_
2049 >>> pd.array([1, 1, pd.NA]).all()
2050 np.True_
2051 >>> pd.array([True, False, pd.NA]).all()
2052 np.False_
2053 >>> pd.array([], dtype="boolean").all()
2054 np.True_
2055 >>> pd.array([pd.NA], dtype="boolean").all()
2056 np.True_
2057 >>> pd.array([pd.NA], dtype="Float64").all()
2058 np.True_
2059
2060 With ``skipna=False``, the result can be NA if this is logically
2061 required (whether ``pd.NA`` is True or False influences the result):
2062
2063 >>> pd.array([True, True, pd.NA]).all(skipna=False)
2064 <NA>
2065 >>> pd.array([1, 1, pd.NA]).all(skipna=False)
2066 <NA>
2067 >>> pd.array([True, False, pd.NA]).all(skipna=False)
2068 np.False_
2069 >>> pd.array([1, 0, pd.NA]).all(skipna=False)
2070 np.False_
2071 """
2072 nv.validate_all((), kwargs)
2073
2074 values = self._data.copy()
2075 np.putmask(values, self._mask, self.dtype._truthy_value)
2076 result = values.all(axis=axis)
2077
2078 if isinstance(result, np.ndarray):
2079 if skipna:
2080 mask = np.zeros(result.shape, dtype=bool)
2081 else:
2082 # Kleene logic: True & NA = NA, False & NA = False
2083 mask = result & self._mask.any(axis=axis)
2084 return self._maybe_mask_result(result, mask)
2085
2086 if skipna:
2087 return result
2088 elif not result or len(self) == 0 or not self._mask.any():
2089 return result
2090 else:
2091 return self.dtype.na_value
2092
2093 def interpolate(
2094 self,
2095 *,
2096 method: InterpolateOptions,
2097 axis: int,
2098 index,
2099 limit,
2100 limit_direction,
2101 limit_area,
2102 copy: bool,
2103 **kwargs,
2104 ) -> FloatingArray:
2105 """
2106 See NDFrame.interpolate.__doc__.
2107 """
2108 # NB: we return type(self) even if copy=False
2109 if self.dtype.kind == "f":
2110 if copy:
2111 data = self._data.copy()
2112 mask = self._mask.copy()
2113 else:
2114 data = self._data
2115 mask = self._mask
2116 elif self.dtype.kind in "iu":
2117 copy = True
2118 data = self._data.astype("f8")
2119 mask = self._mask.copy()
2120 else:
2121 raise NotImplementedError(
2122 f"interpolate is not implemented for dtype={self.dtype}"
2123 )
2124
2125 missing.interpolate_2d_inplace(
2126 data,
2127 method=method,
2128 axis=0,
2129 index=index,
2130 limit=limit,
2131 limit_direction=limit_direction,
2132 limit_area=limit_area,
2133 mask=mask,
2134 **kwargs,
2135 )
2136 if not copy:
2137 return self # type: ignore[return-value]
2138 if self.dtype.kind == "f":
2139 return type(self)._simple_new(data, mask) # type: ignore[return-value]
2140 else:
2141 from pandas.core.arrays import FloatingArray
2142
2143 return FloatingArray._simple_new(data, mask)
2144
2145 def _accumulate(
2146 self, name: str, *, skipna: bool = True, **kwargs
2147 ) -> BaseMaskedArray:
2148 data = self._data
2149 mask = self._mask
2150
2151 op = getattr(masked_accumulations, name)
2152 data, mask = op(data, mask, skipna=skipna, **kwargs)
2153
2154 return self._simple_new(data, mask)
2155
2156 # ------------------------------------------------------------------
2157 # GroupBy Methods
2158
2159 def _groupby_op(
2160 self,
2161 *,
2162 how: str,
2163 has_dropped_na: bool,
2164 min_count: int,
2165 ngroups: int,
2166 ids: npt.NDArray[np.intp],
2167 **kwargs,
2168 ):
2169 from pandas.core.groupby.ops import WrappedCythonOp
2170
2171 op = WrappedCythonOp(how=how, has_dropped_na=has_dropped_na)
2172
2173 # libgroupby functions are responsible for NOT altering mask
2174 mask = self._mask
2175 if op.kind != "aggregate":
2176 result_mask = mask.copy()
2177 else:
2178 result_mask = np.zeros(ngroups, dtype=bool)
2179
2180 if how == "rank" and kwargs.get("na_option") in ["top", "bottom"]:
2181 result_mask[:] = False
2182
2183 res_values = op._cython_op_ndim_compat(
2184 self._data,
2185 min_count=min_count,
2186 ngroups=ngroups,
2187 comp_ids=ids,
2188 mask=mask,
2189 result_mask=result_mask,
2190 **kwargs,
2191 )
2192
2193 if op.how == "ohlc":
2194 arity = op._cython_arity.get(op.how, 1)
2195 result_mask = np.tile(result_mask, (arity, 1)).T
2196
2197 if op.how in ["idxmin", "idxmax"]:
2198 # Result values are indexes to take, keep as ndarray
2199 return res_values
2200 else:
2201 # res_values should already have the correct dtype, we just need to
2202 # wrap in a MaskedArray
2203 return self._maybe_mask_result(res_values, result_mask)
2204
2205 def _groupby_quantile(
2206 self,
2207 *,
2208 qs: npt.NDArray[np.float64],
2209 interpolation: Literal["linear", "lower", "higher", "nearest", "midpoint"],
2210 ids: npt.NDArray[np.intp],
2211 ngroups: int,
2212 starts: npt.NDArray[np.int64],
2213 ends: npt.NDArray[np.int64],
2214 ) -> ArrayLike:
2215 mask = self._mask
2216 nqs = len(qs)
2217 result_mask = np.zeros((ngroups, nqs), dtype=np.bool_)
2218
2219 vals = self.to_numpy(dtype=float, na_value=np.nan)
2220
2221 out = np.empty((ngroups, nqs), dtype=np.float64)
2222 libgroupby.group_quantile(
2223 out,
2224 values=vals,
2225 mask=mask, # type: ignore[arg-type]
2226 labels=ids,
2227 qs=qs,
2228 interpolation=interpolation,
2229 starts=starts,
2230 ends=ends,
2231 result_mask=result_mask,
2232 is_datetimelike=False,
2233 )
2234 out = out.ravel("K") # type: ignore[assignment]
2235 result_mask = result_mask.ravel("K") # type: ignore[assignment]
2236
2237 if not (interpolation in {"linear", "midpoint"} and not is_float_dtype(self)):
2238 # For non-interpolating methods on non-float dtypes, cast back
2239 # to the original numpy dtype (e.g. float64 -> int64).
2240 with warnings.catch_warnings():
2241 warnings.filterwarnings("ignore", category=RuntimeWarning)
2242 out = out.astype(self.dtype.numpy_dtype)
2243
2244 return self._maybe_mask_result(out, result_mask)
2245
2246
2247def transpose_homogeneous_masked_arrays(
2248 masked_arrays: Sequence[BaseMaskedArray],
2249) -> list[BaseMaskedArray]:
2250 """Transpose masked arrays in a list, but faster.
2251
2252 Input should be a list of 1-dim masked arrays of equal length and all have the
2253 same dtype. The caller is responsible for ensuring validity of input data.
2254 """
2255 masked_arrays = list(masked_arrays)
2256 dtype = masked_arrays[0].dtype
2257
2258 values = [arr._data.reshape(1, -1) for arr in masked_arrays]
2259 transposed_values = np.concatenate(
2260 values,
2261 axis=0,
2262 out=np.empty(
2263 (len(masked_arrays), len(masked_arrays[0])),
2264 order="F",
2265 dtype=dtype.numpy_dtype,
2266 ),
2267 )
2268
2269 masks = [arr._mask.reshape(1, -1) for arr in masked_arrays]
2270 transposed_masks = np.concatenate(
2271 masks, axis=0, out=np.empty_like(transposed_values, dtype=bool)
2272 )
2273
2274 arr_type = dtype.construct_array_type()
2275 transposed_arrays: list[BaseMaskedArray] = []
2276 for i in range(transposed_values.shape[1]):
2277 transposed_arr = arr_type(transposed_values[:, i], mask=transposed_masks[:, i])
2278 transposed_arrays.append(transposed_arr)
2279
2280 return transposed_arrays