1from __future__ import annotations
2
3import numbers
4from typing import (
5 TYPE_CHECKING,
6 Any,
7 Self,
8)
9
10import numpy as np
11
12from pandas._config import is_nan_na
13
14from pandas._libs import (
15 lib,
16 missing as libmissing,
17)
18from pandas.errors import AbstractMethodError
19from pandas.util._decorators import cache_readonly
20
21from pandas.core.dtypes.common import (
22 is_integer_dtype,
23 is_string_dtype,
24 pandas_dtype,
25)
26
27from pandas.core.arrays.masked import (
28 BaseMaskedArray,
29 BaseMaskedDtype,
30)
31
32if TYPE_CHECKING:
33 from collections.abc import (
34 Callable,
35 Mapping,
36 )
37
38 import pyarrow
39
40 from pandas._typing import (
41 DtypeObj,
42 npt,
43 )
44
45 from pandas.core.dtypes.dtypes import ExtensionDtype
46
47
48class NumericDtype(BaseMaskedDtype):
49 _default_np_dtype: np.dtype
50 _checker: Callable[[Any], bool] # is_foo_dtype
51
52 def __repr__(self) -> str:
53 return f"{self.name}Dtype()"
54
55 @cache_readonly
56 def is_signed_integer(self) -> bool:
57 return self.kind == "i"
58
59 @cache_readonly
60 def is_unsigned_integer(self) -> bool:
61 return self.kind == "u"
62
63 @property
64 def _is_numeric(self) -> bool:
65 return True
66
67 def __from_arrow__(
68 self, array: pyarrow.Array | pyarrow.ChunkedArray
69 ) -> BaseMaskedArray:
70 """
71 Construct IntegerArray/FloatingArray from pyarrow Array/ChunkedArray.
72 """
73 import pyarrow
74
75 from pandas.core.arrays.arrow._arrow_utils import (
76 pyarrow_array_to_numpy_and_mask,
77 )
78
79 array_class = self.construct_array_type()
80
81 pyarrow_type = pyarrow.from_numpy_dtype(self.type)
82 if not array.type.equals(pyarrow_type) and not pyarrow.types.is_null(
83 array.type
84 ):
85 # test_from_arrow_type_error raise for string, but allow
86 # through itemsize conversion GH#31896
87 rt_dtype = pandas_dtype(array.type.to_pandas_dtype())
88 if rt_dtype.kind not in "iuf":
89 # Could allow "c" or potentially disallow float<->int conversion,
90 # but at the moment we specifically test that uint<->int works
91 raise TypeError(
92 f"Expected array of {self} type, got {array.type} instead"
93 )
94
95 array = array.cast(pyarrow_type)
96
97 if isinstance(array, pyarrow.ChunkedArray):
98 array = array.combine_chunks()
99
100 data, mask = pyarrow_array_to_numpy_and_mask(array, dtype=self.numpy_dtype)
101 if data.dtype.kind == "f" and is_nan_na():
102 mask[np.isnan(data)] = False
103 return array_class(data.copy(), ~mask, copy=False)
104
105 @classmethod
106 def _get_dtype_mapping(cls) -> Mapping[np.dtype, NumericDtype]:
107 raise AbstractMethodError(cls)
108
109 @classmethod
110 def _standardize_dtype(cls, dtype: NumericDtype | str | np.dtype) -> NumericDtype:
111 """
112 Convert a string representation or a numpy dtype to NumericDtype.
113 """
114 if isinstance(dtype, str) and (dtype.startswith(("Int", "UInt", "Float"))):
115 # Avoid DeprecationWarning from NumPy about np.dtype("Int64")
116 # https://github.com/numpy/numpy/pull/7476
117 dtype = dtype.lower()
118
119 if not isinstance(dtype, NumericDtype):
120 mapping = cls._get_dtype_mapping()
121 try:
122 dtype = mapping[np.dtype(dtype)]
123 except KeyError as err:
124 raise ValueError(f"invalid dtype specified {dtype}") from err
125 return dtype
126
127 @classmethod
128 def _safe_cast(cls, values: np.ndarray, dtype: np.dtype, copy: bool) -> np.ndarray:
129 """
130 Safely cast the values to the given dtype.
131
132 "safe" in this context means the casting is lossless.
133 """
134 raise AbstractMethodError(cls)
135
136
137def _coerce_to_data_and_mask(values, dtype, copy: bool, dtype_cls: type[NumericDtype]):
138 checker = dtype_cls._checker
139 default_dtype = dtype_cls._default_np_dtype
140
141 mask = None
142 inferred_type = None
143
144 if dtype is None and hasattr(values, "dtype"):
145 if checker(values.dtype):
146 dtype = values.dtype
147
148 if dtype is not None:
149 dtype = dtype_cls._standardize_dtype(dtype)
150
151 cls = dtype_cls().construct_array_type()
152 if isinstance(values, cls):
153 values, mask = values._data, values._mask
154 if dtype is not None:
155 values = values.astype(dtype.numpy_dtype, copy=False)
156
157 if copy:
158 values = values.copy()
159 mask = mask.copy()
160 return values, mask
161
162 original = values
163 if not copy:
164 values = np.asarray(values)
165 else:
166 values = np.array(values, copy=copy)
167 inferred_type = None
168 if values.dtype == object or is_string_dtype(values.dtype):
169 inferred_type = lib.infer_dtype(values, skipna=True)
170 if inferred_type == "boolean" and dtype is None:
171 # object dtype array of bools
172 name = dtype_cls.__name__.strip("_")
173 raise TypeError(f"{values.dtype} cannot be converted to {name}")
174
175 elif values.dtype.kind == "b" and checker(dtype):
176 # fastpath
177 mask = np.zeros(len(values), dtype=np.bool_)
178 if not copy:
179 values = np.asarray(values, dtype=default_dtype)
180 else:
181 values = np.array(values, dtype=default_dtype, copy=copy)
182
183 elif values.dtype.kind not in "iuf":
184 name = dtype_cls.__name__.strip("_")
185 raise TypeError(f"{values.dtype} cannot be converted to {name}")
186
187 if values.ndim != 1:
188 raise TypeError("values must be a 1D list-like")
189
190 if mask is None:
191 if values.dtype.kind in "iu":
192 # fastpath
193 mask = np.zeros(len(values), dtype=np.bool_)
194 elif values.dtype.kind == "f":
195 # np.isnan is faster than is_numeric_na() for floats
196 # github issue: #60066
197 if is_nan_na():
198 mask = np.isnan(values)
199 else:
200 mask = np.zeros(len(values), dtype=np.bool_)
201 if dtype_cls.__name__.strip("_").startswith(("I", "U")):
202 wrong = np.isnan(values)
203 if wrong.any():
204 raise ValueError("Cannot cast NaN value to Integer dtype.")
205 elif is_nan_na():
206 mask = libmissing.is_numeric_na(values)
207 else:
208 # is_numeric_na will raise on non-numeric NAs
209 libmissing.is_numeric_na(values)
210 mask = libmissing.is_pdna_or_none(values)
211 else:
212 assert len(mask) == len(values)
213
214 if mask.ndim != 1:
215 raise TypeError("mask must be a 1D list-like")
216
217 # infer dtype if needed
218 if dtype is None:
219 dtype = default_dtype
220 else:
221 dtype = dtype.numpy_dtype
222
223 if is_integer_dtype(dtype) and values.dtype.kind == "f" and len(values) > 0:
224 if mask.all():
225 values = np.ones(values.shape, dtype=dtype)
226 else:
227 idx = np.nanargmax(values)
228 if int(values[idx]) != original[idx]:
229 # We have ints that lost precision during the cast.
230 inferred_type = lib.infer_dtype(original, skipna=True)
231 if (
232 inferred_type not in ["floating", "mixed-integer-float"]
233 and not mask.any()
234 ):
235 values = np.asarray(original, dtype=dtype)
236 else:
237 values = np.asarray(original, dtype="object")
238
239 # we copy as need to coerce here
240 if mask.any():
241 values = values.copy()
242 values[mask] = dtype_cls._internal_fill_value
243 if inferred_type in ("string", "unicode"):
244 # casts from str are always safe since they raise
245 # a ValueError if the str cannot be parsed into a float
246 values = values.astype(dtype, copy=copy)
247 else:
248 values = dtype_cls._safe_cast(values, dtype, copy=False)
249 return values, mask
250
251
252class NumericArray(BaseMaskedArray):
253 """
254 Base class for IntegerArray and FloatingArray.
255 """
256
257 _dtype_cls: type[NumericDtype]
258
259 def __init__(
260 self, values: np.ndarray, mask: npt.NDArray[np.bool_], copy: bool = False
261 ) -> None:
262 checker = self._dtype_cls._checker
263 if not (isinstance(values, np.ndarray) and checker(values.dtype)):
264 descr = (
265 "floating"
266 if self._dtype_cls.kind == "f" # type: ignore[comparison-overlap]
267 else "integer"
268 )
269 raise TypeError(
270 f"values should be {descr} numpy array. Use "
271 "the 'pd.array' function instead"
272 )
273 if values.dtype == np.float16:
274 # If we don't raise here, then accessing self.dtype would raise
275 raise TypeError("FloatingArray does not support np.float16 dtype.")
276
277 # NB: if is_nan_na() is True
278 # then caller is responsible for ensuring
279 # assert mask[np.isnan(values)].all()
280
281 super().__init__(values, mask, copy=copy)
282
283 @cache_readonly
284 def dtype(self) -> NumericDtype:
285 mapping = self._dtype_cls._get_dtype_mapping()
286 return mapping[self._data.dtype]
287
288 @classmethod
289 def _coerce_to_array(
290 cls, value, *, dtype: DtypeObj, copy: bool = False
291 ) -> tuple[np.ndarray, np.ndarray]:
292 dtype_cls = cls._dtype_cls
293 values, mask = _coerce_to_data_and_mask(value, dtype, copy, dtype_cls)
294 return values, mask
295
296 @classmethod
297 def _from_sequence_of_strings(
298 cls, strings, *, dtype: ExtensionDtype, copy: bool = False
299 ) -> Self:
300 from pandas.core.tools.numeric import to_numeric
301
302 scalars = to_numeric(strings, errors="raise", dtype_backend="numpy_nullable")
303 return cls._from_sequence(scalars, dtype=dtype, copy=copy)
304
305 _HANDLED_TYPES = (np.ndarray, numbers.Number)