1"""
2Generic data algorithms. This module is experimental at the moment and not
3intended for public consumption
4"""
5
6from __future__ import annotations
7
8import decimal
9import operator
10from typing import (
11 TYPE_CHECKING,
12 Literal,
13 TypeVar,
14 cast,
15 overload,
16)
17import warnings
18
19import numpy as np
20
21from pandas._libs import (
22 algos,
23 hashtable as htable,
24 iNaT,
25 lib,
26)
27from pandas._libs.missing import NA
28from pandas.util._decorators import set_module
29from pandas.util._exceptions import find_stack_level
30
31from pandas.core.dtypes.cast import (
32 construct_1d_object_array_from_listlike,
33 np_find_common_type,
34)
35from pandas.core.dtypes.common import (
36 ensure_float64,
37 ensure_object,
38 ensure_platform_int,
39 is_bool_dtype,
40 is_complex_dtype,
41 is_dict_like,
42 is_dtype_equal,
43 is_extension_array_dtype,
44 is_float,
45 is_float_dtype,
46 is_integer,
47 is_integer_dtype,
48 is_list_like,
49 is_object_dtype,
50 is_signed_integer_dtype,
51 needs_i8_conversion,
52)
53from pandas.core.dtypes.concat import concat_compat
54from pandas.core.dtypes.dtypes import (
55 BaseMaskedDtype,
56 CategoricalDtype,
57 ExtensionDtype,
58 NumpyEADtype,
59)
60from pandas.core.dtypes.generic import (
61 ABCExtensionArray,
62 ABCIndex,
63 ABCMultiIndex,
64 ABCNumpyExtensionArray,
65 ABCSeries,
66)
67from pandas.core.dtypes.missing import (
68 isna,
69 na_value_for_dtype,
70)
71
72from pandas.core.array_algos.take import take_nd
73from pandas.core.construction import (
74 array as pd_array,
75 ensure_wrapped_if_datetimelike,
76 extract_array,
77)
78from pandas.core.indexers import validate_indices
79
80if TYPE_CHECKING:
81 from pandas._typing import (
82 AnyArrayLike,
83 ArrayLike,
84 ArrayLikeT,
85 AxisInt,
86 DtypeObj,
87 ListLike,
88 NumpySorter,
89 NumpyValueArrayLike,
90 RankMethod,
91 RankNaOption,
92 TakeIndexer,
93 npt,
94 )
95
96 from pandas import (
97 Categorical,
98 Index,
99 Series,
100 )
101 from pandas.core.arrays import (
102 BaseMaskedArray,
103 ExtensionArray,
104 )
105
106 T = TypeVar("T", bound=Index | Categorical | ExtensionArray)
107
108
109# --------------- #
110# dtype access #
111# --------------- #
112def _ensure_data(values: ArrayLike) -> np.ndarray:
113 """
114 routine to ensure that our data is of the correct
115 input dtype for lower-level routines
116
117 This will coerce:
118 - ints -> int64
119 - uint -> uint64
120 - bool -> uint8
121 - datetimelike -> i8
122 - datetime64tz -> i8 (in local tz)
123 - categorical -> codes
124
125 Parameters
126 ----------
127 values : np.ndarray or ExtensionArray
128
129 Returns
130 -------
131 np.ndarray
132 """
133
134 if not isinstance(values, ABCMultiIndex):
135 # extract_array would raise
136 values = extract_array(values, extract_numpy=True)
137
138 if is_object_dtype(values.dtype):
139 return ensure_object(np.asarray(values))
140
141 elif isinstance(values.dtype, BaseMaskedDtype):
142 # i.e. BooleanArray, FloatingArray, IntegerArray
143 values = cast("BaseMaskedArray", values)
144 if not values._hasna:
145 # No pd.NAs -> We can avoid an object-dtype cast (and copy) GH#41816
146 # recurse to avoid re-implementing logic for eg bool->uint8
147 return _ensure_data(values._data)
148 return np.asarray(values)
149
150 elif isinstance(values.dtype, CategoricalDtype):
151 # NB: cases that go through here should NOT be using _reconstruct_data
152 # on the back-end.
153 values = cast("Categorical", values)
154 return values.codes
155
156 elif is_bool_dtype(values.dtype):
157 if isinstance(values, np.ndarray):
158 # i.e. actually dtype == np.dtype("bool")
159 return np.asarray(values).view("uint8")
160 else:
161 # e.g. Sparse[bool, False], see test_mode_sparse_bool
162 return np.asarray(values).astype("uint8", copy=False)
163
164 elif is_integer_dtype(values.dtype):
165 return np.asarray(values)
166
167 elif is_float_dtype(values.dtype):
168 # itemsize comes off the values, not the dtype: an ExtensionDtype
169 # (e.g. Sparse[float64]) may not have one.
170 # Note: checking `values.dtype == "float128"` raises on Windows and 32bit
171 float_values = np.asarray(values)
172 if float_values.dtype.itemsize in [2, 12, 16]:
173 # we dont (yet) have float128 hashtable support
174 return ensure_float64(float_values)
175 return float_values
176
177 elif is_complex_dtype(values.dtype):
178 # NumpyExtensionArray needs to be unwrapped to the underlying ndarray
179 # so the hashtable fused-type dispatch works (GH#54761)
180 return np.asarray(values)
181
182 # datetimelike
183 elif needs_i8_conversion(values.dtype):
184 npvalues = values.view("i8")
185 npvalues = cast("np.ndarray", npvalues)
186 return npvalues
187
188 # we have failed, return object
189 values = np.asarray(values, dtype=object)
190 return ensure_object(values)
191
192
193def _reconstruct_data(
194 values: ArrayLikeT, dtype: DtypeObj, original: AnyArrayLike
195) -> ArrayLikeT:
196 """
197 reverse of _ensure_data
198
199 Parameters
200 ----------
201 values : np.ndarray or ExtensionArray
202 dtype : np.dtype or ExtensionDtype
203 original : AnyArrayLike
204
205 Returns
206 -------
207 ExtensionArray or np.ndarray
208 """
209 if isinstance(values, ABCExtensionArray) and values.dtype == dtype:
210 # Catch DatetimeArray/TimedeltaArray
211 return values
212
213 if not isinstance(dtype, np.dtype):
214 # i.e. ExtensionDtype; note we have ruled out above the possibility
215 # that values.dtype == dtype
216 cls = dtype.construct_array_type()
217
218 # error: Incompatible return value type
219 # (got "ExtensionArray",
220 # expected "ndarray[tuple[Any, ...], dtype[Any]]")
221 return cls._from_sequence(values, dtype=dtype) # type: ignore[return-value]
222
223 # error: Incompatible return value type
224 # (got "ndarray[tuple[Any, ...], dtype[Any]]",
225 # expected "ExtensionArray")
226 return values.astype(dtype, copy=False) # type: ignore[return-value]
227
228
229def _ensure_arraylike(values, func_name: str) -> ArrayLike:
230 """
231 ensure that we are arraylike if not already
232 """
233 if not isinstance(
234 values,
235 (ABCIndex, ABCSeries, ABCExtensionArray, np.ndarray, ABCNumpyExtensionArray),
236 ):
237 # GH#52986
238 if func_name != "isin-targets":
239 # Make an exception for the comps argument in isin.
240 raise TypeError(
241 f"{func_name} requires a Series, Index, "
242 f"ExtensionArray, np.ndarray or NumpyExtensionArray "
243 f"got {type(values).__name__}."
244 )
245
246 inferred = lib.infer_dtype(values, skipna=False)
247 if inferred in ["mixed", "string", "mixed-integer"]:
248 # "mixed-integer" to ensure we do not cast ["ss", 42] to str GH#22160
249 if isinstance(values, tuple):
250 values = list(values)
251 values = construct_1d_object_array_from_listlike(values)
252 else:
253 values = np.asarray(values)
254 return values
255
256
257_hashtables = {
258 "complex128": htable.Complex128HashTable,
259 "complex64": htable.Complex64HashTable,
260 "float64": htable.Float64HashTable,
261 "float32": htable.Float32HashTable,
262 "uint64": htable.UInt64HashTable,
263 "uint32": htable.UInt32HashTable,
264 "uint16": htable.UInt16HashTable,
265 "uint8": htable.UInt8HashTable,
266 "int64": htable.Int64HashTable,
267 "int32": htable.Int32HashTable,
268 "int16": htable.Int16HashTable,
269 "int8": htable.Int8HashTable,
270 "string": htable.StringHashTable,
271 "object": htable.PyObjectHashTable,
272}
273
274
275def _get_hashtable_algo(
276 values: np.ndarray,
277) -> tuple[type[htable.HashTable], np.ndarray]:
278 """
279 Parameters
280 ----------
281 values : np.ndarray
282
283 Returns
284 -------
285 htable : HashTable subclass
286 values : ndarray
287 """
288 values = _ensure_data(values)
289
290 ndtype = _check_object_for_strings(values)
291 hashtable = _hashtables[ndtype]
292 return hashtable, values
293
294
295def _check_object_for_strings(values: np.ndarray) -> str:
296 """
297 Check if we can use string hashtable instead of object hashtable.
298
299 Parameters
300 ----------
301 values : ndarray
302
303 Returns
304 -------
305 str
306 """
307 ndtype = values.dtype.name
308 if ndtype == "object":
309 # it's cheaper to use a String Hash Table than Object; we infer
310 # including nulls because that is the only difference between
311 # StringHashTable and ObjectHashtable
312 if lib.is_string_array(values, skipna=False):
313 ndtype = "string"
314 return ndtype
315
316
317# --------------- #
318# top-level algos #
319# --------------- #
320
321
322@overload
323def unique(values: T) -> T: ...
324@overload
325def unique(values: np.ndarray | Series) -> np.ndarray: ...
326
327
328@set_module("pandas")
329def unique(values):
330 """
331 Return unique values based on a hash table.
332
333 Uniques are returned in order of appearance. This does NOT sort.
334
335 Significantly faster than numpy.unique for long enough sequences.
336 Includes NA values.
337
338 Parameters
339 ----------
340 values : 1d array-like
341 The input array-like object containing values from which to extract
342 unique values.
343
344 Returns
345 -------
346 Index, ExtensionArray or numpy.ndarray
347 The return type depends on the type of the input:
348
349 * :class:`Index` : when the input is an :class:`Index`, the result is
350 an :class:`Index` (or a subclass such as :class:`DatetimeIndex`).
351 * :class:`ExtensionArray` : when the input has an
352 :class:`ExtensionDtype` (for example :class:`Categorical`,
353 tz-aware ``datetime64``, :class:`IntervalArray`, a masked
354 integer/boolean/float array, an Arrow-backed array, or
355 :class:`NumpyExtensionArray`), the result is an
356 :class:`ExtensionArray` of the same dtype.
357 * :class:`numpy.ndarray` : for any other input (a NumPy-dtype
358 :class:`Series`, a :class:`numpy.ndarray`, or any other 1-D
359 array-like), the result is a :class:`numpy.ndarray`.
360
361 See Also
362 --------
363 Index.unique : Return unique values from an Index.
364 Series.unique : Return unique values of Series object.
365
366 Notes
367 -----
368 When working with object-dtype arrays, boolean and integer values may not
369 be distinguished since ``True == 1`` and ``False == 0`` in Python.
370
371 >>> pd.unique(np.array([True, 1, False, 0], dtype=object))
372 array([True, False], dtype=object)
373
374 Examples
375 --------
376 >>> pd.unique(pd.Series([2, 1, 3, 3]))
377 array([2, 1, 3])
378
379 >>> pd.unique(pd.Series([2] + [1] * 5))
380 array([2, 1])
381
382 >>> pd.unique(pd.Series([pd.Timestamp("20160101"), pd.Timestamp("20160101")]))
383 array(['2016-01-01T00:00:00.000000'], dtype='datetime64[us]')
384
385 >>> pd.unique(
386 ... pd.Series(
387 ... [
388 ... pd.Timestamp("20160101", tz="US/Eastern"),
389 ... pd.Timestamp("20160101", tz="US/Eastern"),
390 ... ],
391 ... dtype="M8[ns, US/Eastern]",
392 ... )
393 ... )
394 <DatetimeArray>
395 ['2016-01-01 00:00:00-05:00']
396 Length: 1, dtype: datetime64[ns, US/Eastern]
397
398 >>> pd.unique(
399 ... pd.Index(
400 ... [
401 ... pd.Timestamp("20160101", tz="US/Eastern"),
402 ... pd.Timestamp("20160101", tz="US/Eastern"),
403 ... ],
404 ... dtype="M8[ns, US/Eastern]",
405 ... )
406 ... )
407 DatetimeIndex(['2016-01-01 00:00:00-05:00'],
408 dtype='datetime64[ns, US/Eastern]',
409 freq=None)
410
411 >>> pd.unique(np.array(list("baabc"), dtype="O"))
412 array(['b', 'a', 'c'], dtype=object)
413
414 An unordered Categorical will return categories in the
415 order of appearance.
416
417 >>> pd.unique(pd.Series(pd.Categorical(list("baabc"))))
418 ['b', 'a', 'c']
419 Categories (3, str): ['a', 'b', 'c']
420
421 >>> pd.unique(pd.Series(pd.Categorical(list("baabc"), categories=list("abc"))))
422 ['b', 'a', 'c']
423 Categories (3, str): ['a', 'b', 'c']
424
425 An ordered Categorical preserves the category ordering.
426
427 >>> pd.unique(
428 ... pd.Series(
429 ... pd.Categorical(list("baabc"), categories=list("abc"), ordered=True)
430 ... )
431 ... )
432 ['b', 'a', 'c']
433 Categories (3, str): ['a' < 'b' < 'c']
434
435 An array of tuples
436
437 >>> pd.unique(pd.Series([("a", "b"), ("b", "a"), ("a", "c"), ("b", "a")]).values)
438 array([('a', 'b'), ('b', 'a'), ('a', 'c')], dtype=object)
439
440 A NumpyExtensionArray of complex
441
442 >>> pd.unique(pd.array([1 + 1j, 2, 3]))
443 <NumpyExtensionArray>
444 [(1+1j), (2+0j), (3+0j)]
445 Length: 3, dtype: complex128
446 """
447 return unique_with_mask(values)
448
449
450def nunique_ints(values: ArrayLike) -> int:
451 """
452 Return the number of unique values for integer array-likes.
453
454 Significantly faster than pandas.unique for long enough sequences.
455 No checks are done to ensure input is integral.
456
457 Parameters
458 ----------
459 values : 1d array-like
460
461 Returns
462 -------
463 int : The number of unique values in ``values``
464 """
465 if len(values) == 0:
466 return 0
467 values = _ensure_data(values)
468 # bincount requires intp
469 result = (np.bincount(values.ravel().astype("intp")) != 0).sum()
470 return result
471
472
473def unique_with_mask(values, mask: npt.NDArray[np.bool_] | None = None):
474 """See algorithms.unique for docs. Takes a mask for masked arrays."""
475 values = _ensure_arraylike(values, func_name="unique")
476
477 if isinstance(values.dtype, ExtensionDtype):
478 # Dispatch to extension dtype's unique.
479 return values.unique()
480
481 if isinstance(values, ABCIndex):
482 # Dispatch to Index's unique.
483 return values.unique()
484
485 original = values
486 hashtable, values = _get_hashtable_algo(values)
487
488 table = hashtable(len(values))
489 if mask is None:
490 uniques = table.unique(values)
491 uniques = _reconstruct_data(uniques, original.dtype, original)
492 return uniques
493
494 else:
495 uniques, mask = table.unique(values, mask=mask)
496 uniques = _reconstruct_data(uniques, original.dtype, original)
497 assert mask is not None # for mypy
498 return uniques, mask.astype("bool")
499
500
501unique1d = unique
502
503
504_MINIMUM_COMP_ARR_LEN = 1_000_000
505
506# Integers with magnitude <= 2**53 are exactly representable as float64; above
507# this a 64-bit int cast to float64 may collide with a distinct value.
508_FLOAT64_INT_EXACT_MAX = 2**53
509
510
511def _may_lose_precision_as_float64(arr: np.ndarray) -> bool:
512 """
513 Whether casting ``arr`` to float64 could change an integer's identity.
514
515 Only 64-bit integers can carry a magnitude that float64 cannot represent
516 exactly; narrower widths always round-trip.
517 """
518 if arr.dtype.kind not in "iu" or arr.dtype.itemsize < 8 or arr.size == 0:
519 return False
520 return (
521 int(arr.min()) < -_FLOAT64_INT_EXACT_MAX
522 or int(arr.max()) > _FLOAT64_INT_EXACT_MAX
523 )
524
525
526def isin(comps: ListLike, values: ListLike) -> npt.NDArray[np.bool_]:
527 """
528 Compute the isin boolean array.
529
530 Parameters
531 ----------
532 comps : list-like
533 values : list-like
534
535 Returns
536 -------
537 ndarray[bool]
538 Same length as `comps`.
539 """
540 if not is_list_like(comps):
541 raise TypeError(
542 "only list-like objects are allowed to be passed "
543 f"to isin(), you passed a `{type(comps).__name__}`"
544 )
545 if not is_list_like(values):
546 raise TypeError(
547 "only list-like objects are allowed to be passed "
548 f"to isin(), you passed a `{type(values).__name__}`"
549 )
550
551 if isinstance(values, (set, frozenset)) and len(values) > 0:
552 # GH#25507: for a set of values, membership can be tested directly
553 # via the set, avoiding an O(len(values)) materialization that
554 # otherwise dominates when comps is much smaller than values.
555 # Restrict to integer/bool comps (i.e. dtypes that cannot contain
556 # NaN), since Python set membership would mis-handle the case where
557 # both sides contain NaN values that are not identical.
558 if isinstance(comps, (ABCSeries, ABCIndex)):
559 comps_arr = comps._values
560 else:
561 comps_arr = comps
562 if (
563 isinstance(comps_arr, np.ndarray)
564 and comps_arr.ndim == 1
565 and comps_arr.dtype.kind in "iub"
566 # Only worth it when comps is smaller than values: this path does
567 # an O(len(comps)) Python-level scan, which is much slower per
568 # element than the vectorized htable path below. It wins only by
569 # skipping the O(len(values)) materialization of the set, so it
570 # must not fire for the common large-comps/small-set case.
571 and comps_arr.size < len(values)
572 # A 64-bit comp with magnitude >= 2**53 can match a float in the
573 # set lossily under the materialized path below (via a float64
574 # cast) but never under exact Python equality, so the result would
575 # depend on which path is taken. Only use this path when comps
576 # values are strictly inside float64's exact-int range, where the
577 # two agree.
578 and (
579 comps_arr.dtype.itemsize < 8
580 or comps_arr.size == 0
581 or (
582 -_FLOAT64_INT_EXACT_MAX < int(comps_arr.min())
583 and int(comps_arr.max()) < _FLOAT64_INT_EXACT_MAX
584 )
585 )
586 ):
587 return np.fromiter(
588 (item in values for item in comps_arr.tolist()),
589 dtype=bool,
590 count=comps_arr.size,
591 )
592
593 if not isinstance(values, (ABCIndex, ABCSeries, ABCExtensionArray, np.ndarray)):
594 orig_values = list(values)
595 values = _ensure_arraylike(orig_values, func_name="isin-targets")
596
597 if (
598 len(values) > 0
599 and values.dtype.kind in "iufcb"
600 and not is_signed_integer_dtype(comps)
601 and not is_dtype_equal(values, comps)
602 ):
603 # GH#46485: casting a 64-bit integer to float64 loses precision
604 # for magnitudes > 2**53, so a numeric comparison can give wrong
605 # results. Recast to object only when such a lossy cast is actually
606 # at risk; otherwise keep the numeric ndarray so the downstream
607 # htable dispatch can use the fast numeric path.
608 # TODO: Share with _find_common_type_compat
609 comps_dtype = getattr(comps, "dtype", None)
610 # values came from _ensure_arraylike's numeric branch, so its
611 # dtype is an np.dtype.
612 values_dtype = cast("np.dtype", values.dtype)
613 if not isinstance(comps_dtype, np.dtype):
614 needs_object = True
615 else:
616 common = np_find_common_type(values_dtype, comps_dtype)
617 if common.kind not in "iufcb":
618 needs_object = True
619 elif common.kind not in "fc":
620 # integer/bool common type -> no lossy float cast
621 # (complex128 components are float64, so "c" is as lossy
622 # as "f")
623 needs_object = False
624 elif comps_dtype.kind in "iu" and comps_dtype.itemsize == 8:
625 # comps is the (potentially large) caller and a 64-bit
626 # integer that would be cast down to float64; too costly to
627 # inspect element-wise, so fall back to object to be safe.
628 needs_object = True
629 elif values_dtype.kind in "iu" and values_dtype.itemsize == 8:
630 # values is the (small) targets list -> cheap to check
631 # exactly. Only object-cast if a target falls outside
632 # float64's exact integer range; otherwise the numeric
633 # cast is lossless and the fast path stays correct.
634 # values is an ndarray here (it came from the numeric
635 # branch of _ensure_arraylike above).
636 values_arr = cast("np.ndarray", values)
637 needs_object = _may_lose_precision_as_float64(values_arr)
638 else:
639 # float/complex values_dtype means _ensure_arraylike may
640 # already have rounded large ints in a mixed int/float
641 # list, so exactness cannot be checked from the array;
642 # recover it from orig_values via the object cast. Smaller
643 # ints/bools cast to float64 exactly.
644 needs_object = values_dtype.kind in "fc"
645 if needs_object:
646 values = construct_1d_object_array_from_listlike(orig_values)
647
648 elif isinstance(values, ABCMultiIndex):
649 # Avoid raising in extract_array
650 values = np.array(values)
651 else:
652 values = extract_array(values, extract_numpy=True, extract_range=True)
653
654 comps_array = _ensure_arraylike(comps, func_name="isin")
655 comps_array = extract_array(comps_array, extract_numpy=True)
656 if not isinstance(comps_array, np.ndarray):
657 # i.e. Extension Array
658 return comps_array.isin(values)
659
660 elif needs_i8_conversion(comps_array.dtype):
661 # Dispatch to DatetimeLikeArrayMixin.isin
662 return pd_array(comps_array).isin(values)
663 elif needs_i8_conversion(values.dtype) and not is_object_dtype(comps_array.dtype):
664 # e.g. comps_array are integers and values are datetime64s
665 return np.zeros(comps_array.shape, dtype=bool)
666 # TODO: not quite right ... Sparse/Categorical
667 elif needs_i8_conversion(values.dtype):
668 return isin(comps_array, values.astype(object))
669
670 elif isinstance(values.dtype, ExtensionDtype):
671 return isin(np.asarray(comps_array), np.asarray(values))
672
673 # GH#59609: uint64 mixed with int64 has float64 as its common dtype, which
674 # is lossy above 2**53. Targets outside the other side's range can never
675 # match, so dropping them lets the rest be cast losslessly.
676 if comps_array.dtype == np.uint64 and values.dtype == np.int64:
677 values = values[values >= 0].astype(np.uint64)
678 elif comps_array.dtype == np.int64 and values.dtype == np.uint64:
679 values = values[values <= np.iinfo(np.int64).max].astype(np.int64)
680
681 # GH#46485, GH#59609: whatever remains may still have a float64/complex128
682 # common dtype, which cannot represent every 64-bit integer; comparing
683 # through it would report a match between distinct values above 2**53.
684 common = np_find_common_type(values.dtype, comps_array.dtype)
685 lossy_common = common.kind in "fc" and (
686 _may_lose_precision_as_float64(values)
687 or _may_lose_precision_as_float64(comps_array)
688 )
689
690 # GH16012
691 # Ensure np.isin doesn't get object types or it *may* throw an exception
692 # Albeit hashmap has O(1) look-up (vs. O(logn) in sorted array),
693 # isin is faster for small sizes
694
695 # GH60678
696 # Ensure values don't contain <NA>, otherwise it throws exception with np.in1d
697
698 if (
699 len(comps_array) > _MINIMUM_COMP_ARR_LEN
700 and len(values) <= 26
701 and comps_array.dtype != object
702 # np.isin would do the same lossy promotion internally
703 and not lossy_common
704 and not any(v is NA for v in values)
705 ):
706 # If the values include nan we need to check for nan explicitly
707 # since np.nan it not equal to np.nan
708 if isna(values).any():
709
710 def f(c, v):
711 return np.logical_or(np.isin(c, v).ravel(), np.isnan(c))
712
713 else:
714 f = lambda a, b: np.isin(a, b).ravel()
715
716 else:
717 if lossy_common:
718 # Compare as Python ints instead of through the lossy float.
719 common = np.dtype(object)
720 values = values.astype(common, copy=False)
721 comps_array = comps_array.astype(common, copy=False)
722 f = htable.ismember
723
724 return f(comps_array, values)
725
726
727def factorize_array(
728 values: np.ndarray,
729 use_na_sentinel: bool = True,
730 size_hint: int | None = None,
731 na_value: object = None,
732 mask: npt.NDArray[np.bool_] | None = None,
733) -> tuple[npt.NDArray[np.intp], np.ndarray]:
734 """
735 Factorize a numpy array to codes and uniques.
736
737 This doesn't do any coercion of types or unboxing before factorization.
738
739 Parameters
740 ----------
741 values : ndarray
742 use_na_sentinel : bool, default True
743 If True, the sentinel -1 will be used for NaN values. If False,
744 NaN values will be encoded as non-negative integers and will not drop the
745 NaN from the uniques of the values.
746 size_hint : int, optional
747 Passed through to the hashtable's 'get_labels' method
748 na_value : object, optional
749 A value in `values` to consider missing. Note: only use this
750 parameter when you know that you don't have any values pandas would
751 consider missing in the array (NaN for float data, iNaT for
752 datetimes, etc.).
753 mask : ndarray[bool], optional
754 If not None, the mask is used as indicator for missing values
755 (True = missing, False = valid) instead of `na_value` or
756 condition "val != val".
757
758 Returns
759 -------
760 codes : ndarray[np.intp]
761 uniques : ndarray
762 """
763 original = values
764 if values.dtype.kind in "mM":
765 # _get_hashtable_algo will cast dt64/td64 to i8 via _ensure_data, so we
766 # need to do the same to na_value. We are assuming here that the passed
767 # na_value is an appropriately-typed NaT.
768 # e.g. test_where_datetimelike_categorical
769 na_value = iNaT
770
771 hash_klass, values = _get_hashtable_algo(values)
772
773 table = hash_klass(size_hint or len(values))
774 uniques, codes = table.factorize(
775 values,
776 na_sentinel=-1,
777 na_value=na_value,
778 mask=mask,
779 ignore_na=use_na_sentinel,
780 )
781
782 # re-cast e.g. i8->dt64/td64, uint8->bool
783 uniques = _reconstruct_data(uniques, original.dtype, original)
784
785 codes = ensure_platform_int(codes)
786 return codes, uniques
787
788
789def factorize_monotonic_codes(
790 arr: np.ndarray,
791 ascending: bool,
792 sort: bool,
793) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]:
794 """
795 Factorize an array known to be monotonic and of length >= 2.
796
797 Uses adjacent-element comparison instead of hash table construction.
798 The caller is responsible for having verified monotonicity; monotonic
799 arrays contain no NA values (NAs break monotonicity checks), so NA
800 handling is not needed.
801
802 Parameters
803 ----------
804 arr : ndarray
805 ascending : bool
806 Whether `arr` is monotonic increasing (otherwise decreasing).
807 sort : bool
808 Whether the uniques should be in sorted order.
809
810 Returns
811 -------
812 codes : ndarray[np.intp]
813 uniques_indexer : ndarray[np.intp]
814 Positions in `arr` of the unique values, ordered to match `codes`.
815 """
816 n = len(arr)
817
818 # Find group transitions via adjacent-element comparison.
819 transitions = np.empty(n, dtype=bool)
820 transitions[0] = True
821 transitions[1:] = arr[1:] != arr[:-1]
822
823 # Build codes using repeat (faster than cumsum on bool).
824 group_starts = np.flatnonzero(transitions)
825 ngroups = len(group_starts)
826 run_lengths = np.diff(group_starts, append=n)
827 codes = np.repeat(np.arange(ngroups, dtype=np.intp), run_lengths)
828
829 if sort and not ascending:
830 # Reverse uniques to sorted order and remap codes.
831 codes = ensure_platform_int(ngroups - 1 - codes)
832 group_starts = group_starts[::-1]
833
834 return codes, group_starts
835
836
837@set_module("pandas")
838def factorize(
839 values,
840 sort: bool = False,
841 use_na_sentinel: bool = True,
842 size_hint: int | None = None,
843) -> tuple[np.ndarray, np.ndarray | Index]:
844 """
845 Encode the object as an enumerated type or categorical variable.
846
847 This method is useful for obtaining a numeric representation of an
848 array when all that matters is identifying distinct values. `factorize`
849 is available as both a top-level function :func:`pandas.factorize`,
850 and as a method :meth:`Series.factorize` and :meth:`Index.factorize`.
851
852 Parameters
853 ----------
854 values : sequence
855 A 1-D sequence. Sequences that aren't pandas objects are
856 coerced to ndarrays before factorization.
857 sort : bool, default False
858 Sort `uniques` and shuffle `codes` to maintain the
859 relationship.
860 use_na_sentinel : bool, default True
861 If True, the sentinel -1 will be used for NaN values. If False,
862 NaN values will be encoded as non-negative integers and will not drop the
863 NaN from the uniques of the values.
864 size_hint : int, optional
865 Hint to the hashtable sizer.
866
867 Returns
868 -------
869 codes : ndarray
870 An integer ndarray that's an indexer into `uniques`.
871 ``uniques.take(codes)`` will have the same values as `values`.
872 uniques : ndarray, Index, or Categorical
873 The unique valid values. When `values` is Categorical, `uniques`
874 is a Categorical. When `values` is some other pandas object, an
875 `Index` is returned. Otherwise, a 1-D ndarray is returned.
876
877 .. note::
878
879 Even if there's a missing value in `values`, `uniques` will
880 *not* contain an entry for it.
881
882 See Also
883 --------
884 cut : Discretize continuous-valued array.
885 unique : Find the unique value in an array.
886
887 Notes
888 -----
889 Reference :ref:`the user guide <reshaping.factorize>` for more examples.
890
891 When working with object-dtype arrays, boolean and integer values may not
892 be distinguished since ``True == 1`` and ``False == 0`` in Python.
893
894 >>> codes, uniques = pd.factorize(np.array([True, 1, False, 0], dtype=object))
895 >>> codes
896 array([0, 0, 1, 1])
897 >>> uniques
898 array([True, False], dtype=object)
899
900 Examples
901 --------
902 These examples all show factorize as a top-level method like
903 ``pd.factorize(values)``. The results are identical for methods like
904 :meth:`Series.factorize`.
905
906 >>> codes, uniques = pd.factorize(np.array(["b", "b", "a", "c", "b"], dtype="O"))
907 >>> codes
908 array([0, 0, 1, 2, 0])
909 >>> uniques
910 array(['b', 'a', 'c'], dtype=object)
911
912 With ``sort=True``, the `uniques` will be sorted, and `codes` will be
913 shuffled so that the relationship is the maintained.
914
915 >>> codes, uniques = pd.factorize(
916 ... np.array(["b", "b", "a", "c", "b"], dtype="O"), sort=True
917 ... )
918 >>> codes
919 array([1, 1, 0, 2, 1])
920 >>> uniques
921 array(['a', 'b', 'c'], dtype=object)
922
923 When ``use_na_sentinel=True`` (the default), missing values are indicated in
924 the `codes` with the sentinel value ``-1`` and missing values are not
925 included in `uniques`.
926
927 >>> codes, uniques = pd.factorize(np.array(["b", None, "a", "c", "b"], dtype="O"))
928 >>> codes
929 array([ 0, -1, 1, 2, 0])
930 >>> uniques
931 array(['b', 'a', 'c'], dtype=object)
932
933 Thus far, we've only factorized lists (which are internally coerced to
934 NumPy arrays). When factorizing pandas objects, the type of `uniques`
935 will differ. For Categoricals, a `Categorical` is returned.
936
937 >>> cat = pd.Categorical(["a", "a", "c"], categories=["a", "b", "c"])
938 >>> codes, uniques = pd.factorize(cat)
939 >>> codes
940 array([0, 0, 1])
941 >>> uniques
942 ['a', 'c']
943 Categories (3, str): ['a', 'b', 'c']
944
945 Notice that ``'b'`` is in ``uniques.categories``, despite not being
946 present in ``cat.values``.
947
948 For all other pandas objects, an Index of the appropriate type is
949 returned.
950
951 >>> cat = pd.Series(["a", "a", "c"])
952 >>> codes, uniques = pd.factorize(cat)
953 >>> codes
954 array([0, 0, 1])
955 >>> uniques
956 Index(['a', 'c'], dtype='str')
957
958 If NaN is in the values, and we want to include NaN in the uniques of the
959 values, it can be achieved by setting ``use_na_sentinel=False``.
960
961 >>> values = np.array([1, 2, 1, np.nan])
962 >>> codes, uniques = pd.factorize(values) # default: use_na_sentinel=True
963 >>> codes
964 array([ 0, 1, 0, -1])
965 >>> uniques
966 array([1., 2.])
967
968 >>> codes, uniques = pd.factorize(values, use_na_sentinel=False)
969 >>> codes
970 array([0, 1, 0, 2])
971 >>> uniques
972 array([ 1., 2., nan])
973 """
974 # Implementation notes: This method is responsible for 3 things
975 # 1.) coercing data to array-like (ndarray, Index, extension array)
976 # 2.) factorizing codes and uniques
977 # 3.) Maybe boxing the uniques in an Index
978 #
979 # Step 2 is dispatched to extension types (like Categorical). They are
980 # responsible only for factorization. All data coercion, sorting and boxing
981 # should happen here.
982 if isinstance(values, (ABCIndex, ABCSeries)):
983 return values.factorize(sort=sort, use_na_sentinel=use_na_sentinel)
984
985 values = _ensure_arraylike(values, func_name="factorize")
986 original = values
987
988 if not isinstance(values, np.ndarray):
989 # i.e. ExtensionArray
990 codes, uniques = values.factorize(use_na_sentinel=use_na_sentinel)
991
992 else:
993 values = np.asarray(values) # convert DTA/TDA/MultiIndex
994
995 if not use_na_sentinel and values.dtype == object:
996 # factorize can now handle differentiating various types of null values.
997 # These can only occur when the array has object dtype.
998 # However, for backwards compatibility we only use the null for the
999 # provided dtype. This may be revisited in the future, see GH#48476.
1000 null_mask = isna(values)
1001 if null_mask.any():
1002 na_value = na_value_for_dtype(values.dtype, compat=False)
1003 # Don't modify (potentially user-provided) array
1004 # error: Argument 2 to "where" has incompatible type "Scalar";
1005 # expected "_Buffer | _SupportsArray[dtype[Any]] | ..."
1006 values = np.where(null_mask, na_value, values) # type: ignore[arg-type]
1007
1008 codes, uniques = factorize_array(
1009 values,
1010 use_na_sentinel=use_na_sentinel,
1011 size_hint=size_hint,
1012 )
1013
1014 if sort and len(uniques) > 0:
1015 uniques, codes = safe_sort(
1016 uniques,
1017 codes,
1018 use_na_sentinel=use_na_sentinel,
1019 assume_unique=True,
1020 verify=False,
1021 )
1022
1023 uniques = _reconstruct_data(uniques, original.dtype, original)
1024
1025 return codes, uniques
1026
1027
1028def value_counts_internal(
1029 values,
1030 sort: bool = True,
1031 ascending: bool = False,
1032 normalize: bool = False,
1033 bins=None,
1034 dropna: bool = True,
1035) -> Series:
1036 from pandas import (
1037 DatetimeIndex,
1038 Index,
1039 Series,
1040 TimedeltaIndex,
1041 )
1042
1043 index_name = getattr(values, "name", None)
1044 name = "proportion" if normalize else "count"
1045
1046 if bins is not None:
1047 from pandas.core.reshape.tile import cut
1048
1049 if isinstance(values, Series):
1050 values = values._values
1051
1052 try:
1053 ii = cut(values, bins, include_lowest=True)
1054 except TypeError as err:
1055 raise TypeError("bins argument only works with numeric data.") from err
1056
1057 # count, remove nulls (from the index), and but the bins
1058 result = ii.value_counts(dropna=dropna)
1059 result.name = name
1060 result = result[result.index.notna()]
1061 result.index = result.index.astype("interval")
1062 result = result.sort_index()
1063
1064 # if we are dropna and we have NO values
1065 if dropna and (result._values == 0).all():
1066 result = result.iloc[0:0]
1067
1068 # normalizing is by len of all (regardless of dropna)
1069 normalize_denominator = len(ii)
1070
1071 else:
1072 normalize_denominator = None
1073 if is_extension_array_dtype(values):
1074 # handle Categorical and sparse,
1075 result = Series(values, copy=False)._values.value_counts(dropna=dropna)
1076 result.name = name
1077 result.index.name = index_name
1078
1079 elif isinstance(values, ABCMultiIndex):
1080 # GH49558
1081 levels = list(range(values.nlevels))
1082 result = (
1083 Series(index=values, name=name)
1084 .groupby(level=levels, dropna=dropna)
1085 .size()
1086 )
1087 result.index.names = values.names
1088
1089 else:
1090 values = _ensure_arraylike(values, func_name="value_counts")
1091 keys, counts, _ = value_counts_arraylike(values, dropna)
1092 if keys.dtype == np.float16:
1093 keys = keys.astype(np.float32)
1094
1095 # Starting in 3.0, we no longer perform dtype inference on the
1096 # Index object we construct here, xref GH#56161
1097 idx = Index(keys, dtype=keys.dtype, name=index_name, copy=False)
1098
1099 if (
1100 not sort
1101 and isinstance(values, (DatetimeIndex, TimedeltaIndex))
1102 and idx.equals(values)
1103 and values._inferred_freq_str is not None
1104 ):
1105 # Preserve freq of original index
1106 idx.freq = values._inferred_freq_str # type: ignore[attr-defined]
1107
1108 result = Series(counts, index=idx, name=name, copy=False)
1109
1110 if sort:
1111 result = result.sort_values(ascending=ascending, kind="stable")
1112
1113 if normalize:
1114 if normalize_denominator is not None:
1115 result = result / normalize_denominator
1116 else:
1117 result = result / result.sum()
1118
1119 return result
1120
1121
1122# Called once from SparseArray, otherwise could be private
1123def value_counts_arraylike(
1124 values: np.ndarray, dropna: bool, mask: npt.NDArray[np.bool_] | None = None
1125) -> tuple[ArrayLike, npt.NDArray[np.int64], int]:
1126 """
1127 Parameters
1128 ----------
1129 values : np.ndarray
1130 dropna : bool
1131 mask : np.ndarray[bool] or None, default None
1132
1133 Returns
1134 -------
1135 uniques : np.ndarray
1136 counts : np.ndarray[np.int64]
1137 """
1138 original = values
1139 values = _ensure_data(values)
1140
1141 keys, counts, na_counter = htable.value_count(values, dropna, mask=mask)
1142
1143 if needs_i8_conversion(original.dtype):
1144 # datetime, timedelta, or period
1145
1146 if dropna:
1147 mask = keys != iNaT
1148 keys, counts = keys[mask], counts[mask]
1149
1150 res_keys = _reconstruct_data(keys, original.dtype, original)
1151 return res_keys, counts, na_counter
1152
1153
1154def duplicated(
1155 values: ArrayLike,
1156 keep: Literal["first", "last", False] = "first",
1157 mask: npt.NDArray[np.bool_] | None = None,
1158) -> npt.NDArray[np.bool_]:
1159 """
1160 Return boolean ndarray denoting duplicate values.
1161
1162 Parameters
1163 ----------
1164 values : np.ndarray or ExtensionArray
1165 Array over which to check for duplicate values.
1166 keep : {'first', 'last', False}, default 'first'
1167 - ``first`` : Mark duplicates as ``True`` except for the first
1168 occurrence.
1169 - ``last`` : Mark duplicates as ``True`` except for the last
1170 occurrence.
1171 - False : Mark all duplicates as ``True``.
1172 mask : ndarray[bool], optional
1173 array indicating which elements to exclude from checking
1174
1175 Returns
1176 -------
1177 duplicated : ndarray[bool]
1178 """
1179 values = _ensure_data(values)
1180 return htable.duplicated(values, keep=keep, mask=mask)
1181
1182
1183def mode(
1184 values: ArrayLike, dropna: bool = True, mask: npt.NDArray[np.bool_] | None = None
1185) -> tuple[np.ndarray, npt.NDArray[np.bool_]] | ExtensionArray:
1186 """
1187 Returns the mode(s) of an array.
1188
1189 Parameters
1190 ----------
1191 values : array-like
1192 Array over which to check for duplicate values.
1193 dropna : bool, default True
1194 Don't consider counts of NaN/NaT.
1195
1196 Returns
1197 -------
1198 Union[Tuple[np.ndarray, npt.NDArray[np.bool_]], ExtensionArray]
1199 """
1200 values = _ensure_arraylike(values, func_name="mode")
1201 original = values
1202
1203 if needs_i8_conversion(values.dtype):
1204 # Got here with ndarray; dispatch to DatetimeArray/TimedeltaArray.
1205 values = ensure_wrapped_if_datetimelike(values)
1206 values = cast("ExtensionArray", values)
1207 return values._mode(dropna=dropna)
1208
1209 values = _ensure_data(values)
1210
1211 npresult, res_mask = htable.mode(values, dropna=dropna, mask=mask)
1212 if res_mask is None:
1213 res_mask = np.zeros(npresult.shape, dtype=np.bool_)
1214 else:
1215 return npresult, res_mask
1216
1217 try:
1218 npresult = safe_sort(npresult)
1219 except TypeError as err:
1220 warnings.warn(
1221 f"Unable to sort modes: {err}",
1222 stacklevel=find_stack_level(),
1223 )
1224
1225 result = _reconstruct_data(npresult, original.dtype, original)
1226 return result, res_mask
1227
1228
1229def rank(
1230 values: ArrayLike,
1231 axis: AxisInt = 0,
1232 method: RankMethod = "average",
1233 na_option: RankNaOption = "keep",
1234 ascending: bool = True,
1235 pct: bool = False,
1236 mask: npt.NDArray[np.bool_] | None = None,
1237) -> npt.NDArray[np.float64]:
1238 """
1239 Rank the values along a given axis.
1240
1241 Parameters
1242 ----------
1243 values : np.ndarray or ExtensionArray
1244 Array whose values will be ranked. The number of dimensions in this
1245 array must not exceed 2.
1246 axis : int, default 0
1247 Axis over which to perform rankings.
1248 method : {'average', 'min', 'max', 'first', 'dense'}, default 'average'
1249 The method by which tiebreaks are broken during the ranking.
1250 na_option : {'keep', 'top'}, default 'keep'
1251 The method by which NaNs are placed in the ranking.
1252 - ``keep``: rank each NaN value with a NaN ranking
1253 - ``top``: replace each NaN with either +/- inf so that they
1254 there are ranked at the top
1255 ascending : bool, default True
1256 Whether or not the elements should be ranked in ascending order.
1257 pct : bool, default False
1258 Whether or not to the display the returned rankings in integer form
1259 (e.g. 1, 2, 3) or in percentile form (e.g. 0.333..., 0.666..., 1).
1260 mask : bool ndarray, optional
1261 Boolean array indicating which elements to exclude from ranking.
1262 """
1263 is_datetimelike = needs_i8_conversion(values.dtype)
1264 values = _ensure_data(values)
1265
1266 if values.ndim == 1:
1267 ranks = algos.rank_1d(
1268 values,
1269 is_datetimelike=is_datetimelike,
1270 ties_method=method,
1271 ascending=ascending,
1272 na_option=na_option,
1273 pct=pct,
1274 mask=mask,
1275 )
1276 elif values.ndim == 2:
1277 assert mask is None
1278 ranks = algos.rank_2d(
1279 values,
1280 axis=axis,
1281 is_datetimelike=is_datetimelike,
1282 ties_method=method,
1283 ascending=ascending,
1284 na_option=na_option,
1285 pct=pct,
1286 )
1287 else:
1288 raise TypeError("Array with ndim > 2 is not supported.")
1289
1290 return ranks
1291
1292
1293def is_monotonic(values: ArrayLike) -> tuple[bool, bool, bool]:
1294 """
1295 Determine whether values are monotonic increasing/decreasing.
1296
1297 Parameters
1298 ----------
1299 values : np.ndarray or ExtensionArray
1300
1301 Returns
1302 -------
1303 tuple[bool, bool, bool]
1304 (is_monotonic_increasing, is_monotonic_decreasing, is_strict_monotonic)
1305
1306 Raises
1307 ------
1308 TypeError
1309 If the dtype is not orderable by the underlying routine (e.g. complex).
1310 """
1311 timelike = needs_i8_conversion(values.dtype)
1312 values = _ensure_data(values)
1313 inc, dec, strict = algos.is_monotonic(values, timelike=timelike)
1314 return bool(inc), bool(dec), bool(strict)
1315
1316
1317# ---- #
1318# take #
1319# ---- #
1320
1321
1322def occurrence_rank(codes: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]:
1323 """
1324 For each element, the number of earlier elements with the same code.
1325
1326 Parameters
1327 ----------
1328 codes : np.ndarray[intp]
1329
1330 Returns
1331 -------
1332 np.ndarray[intp]
1333
1334 Examples
1335 --------
1336 >>> occurrence_rank(np.array([0, 1, 0, 0, 1]))
1337 array([0, 0, 1, 2, 1])
1338 """
1339 order = np.argsort(codes, kind="stable")
1340 sorted_codes = codes[order]
1341 n = len(codes)
1342 is_start = np.empty(n, dtype=bool)
1343 is_start[:1] = True
1344 is_start[1:] = sorted_codes[1:] != sorted_codes[:-1]
1345 group_start = np.maximum.accumulate(np.where(is_start, np.arange(n), 0))
1346 rank = np.empty(n, dtype=np.intp)
1347 rank[order] = np.arange(n) - group_start
1348 return rank
1349
1350
1351@set_module("pandas.api.extensions")
1352def take(
1353 arr,
1354 indices: TakeIndexer,
1355 axis: AxisInt = 0,
1356 allow_fill: bool = False,
1357 fill_value=None,
1358):
1359 """
1360 Take elements from an array.
1361
1362 Parameters
1363 ----------
1364 arr : numpy.ndarray, ExtensionArray, Index, or Series
1365 Input array.
1366 indices : sequence of int or one-dimensional np.ndarray of int
1367 Indices to be taken.
1368 axis : int, default 0
1369 The axis over which to select values.
1370 allow_fill : bool, default False
1371 How to handle negative values in `indices`.
1372
1373 * False: negative values in `indices` indicate positional indices
1374 from the right (the default). This is similar to :func:`numpy.take`.
1375
1376 * True: negative values in `indices` indicate
1377 missing values. These values are set to `fill_value`. Any other
1378 negative values raise a ``ValueError``.
1379
1380 fill_value : any, optional
1381 Fill value to use for NA-indices when `allow_fill` is True.
1382 This may be ``None``, in which case the default NA value for
1383 the type (``self.dtype.na_value``) is used.
1384
1385 For multi-dimensional `arr`, each *element* is filled with
1386 `fill_value`.
1387
1388 Returns
1389 -------
1390 ndarray or ExtensionArray
1391 Same type as the input.
1392
1393 Raises
1394 ------
1395 IndexError
1396 When `indices` is out of bounds for the array.
1397 ValueError
1398 When the indexer contains negative values other than ``-1``
1399 and `allow_fill` is True.
1400
1401 Notes
1402 -----
1403 When `allow_fill` is False, `indices` may be whatever dimensionality
1404 is accepted by NumPy for `arr`.
1405
1406 When `allow_fill` is True, `indices` should be 1-D.
1407
1408 See Also
1409 --------
1410 numpy.take : Take elements from an array along an axis.
1411
1412 Examples
1413 --------
1414 >>> import pandas as pd
1415
1416 With the default ``allow_fill=False``, negative numbers indicate
1417 positional indices from the right.
1418
1419 >>> pd.api.extensions.take(np.array([10, 20, 30]), [0, 0, -1])
1420 array([10, 10, 30])
1421
1422 Setting ``allow_fill=True`` will place `fill_value` in those positions.
1423
1424 >>> pd.api.extensions.take(np.array([10, 20, 30]), [0, 0, -1], allow_fill=True)
1425 array([10., 10., nan])
1426
1427 >>> pd.api.extensions.take(
1428 ... np.array([10, 20, 30]), [0, 0, -1], allow_fill=True, fill_value=-10
1429 ... )
1430 array([ 10, 10, -10])
1431 """
1432 if not isinstance(
1433 arr,
1434 (np.ndarray, ABCExtensionArray, ABCIndex, ABCSeries, ABCNumpyExtensionArray),
1435 ):
1436 # GH#52981
1437 raise TypeError(
1438 "pd.api.extensions.take requires a numpy.ndarray, ExtensionArray, "
1439 f"Index, Series, or NumpyExtensionArray got {type(arr).__name__}."
1440 )
1441
1442 indices = ensure_platform_int(indices)
1443
1444 if allow_fill:
1445 # Pandas style, -1 means NA
1446 validate_indices(indices, arr.shape[axis])
1447 # error: Argument 1 to "take_nd" has incompatible type
1448 # "ndarray[Any, Any] | ExtensionArray | Index | Series"; expected
1449 # "ndarray[Any, Any]"
1450 result = take_nd(
1451 arr, # type: ignore[arg-type]
1452 indices,
1453 axis=axis,
1454 allow_fill=True,
1455 fill_value=fill_value,
1456 )
1457 else:
1458 # NumPy style
1459 # error: Unexpected keyword argument "axis" for "take" of "ExtensionArray"
1460 result = arr.take(indices, axis=axis) # type: ignore[call-arg,assignment]
1461 return result
1462
1463
1464# ------------ #
1465# searchsorted #
1466# ------------ #
1467
1468
1469def searchsorted(
1470 arr: ArrayLike,
1471 value: NumpyValueArrayLike | ExtensionArray,
1472 side: Literal["left", "right"] = "left",
1473 sorter: NumpySorter | None = None,
1474) -> npt.NDArray[np.intp] | np.intp:
1475 """
1476 Find indices where elements should be inserted to maintain order.
1477
1478 Find the indices into a sorted array `arr` (a) such that, if the
1479 corresponding elements in `value` were inserted before the indices,
1480 the order of `arr` would be preserved.
1481
1482 Assuming that `arr` is sorted:
1483
1484 ====== ================================
1485 `side` returned index `i` satisfies
1486 ====== ================================
1487 left ``arr[i-1] < value <= self[i]``
1488 right ``arr[i-1] <= value < self[i]``
1489 ====== ================================
1490
1491 Parameters
1492 ----------
1493 arr: np.ndarray, ExtensionArray, Series
1494 Input array. If `sorter` is None, then it must be sorted in
1495 ascending order, otherwise `sorter` must be an array of indices
1496 that sort it.
1497 value : array-like or scalar
1498 Values to insert into `arr`.
1499 side : {'left', 'right'}, optional
1500 If 'left', the index of the first suitable location found is given.
1501 If 'right', return the last such index. If there is no suitable
1502 index, return either 0 or N (where N is the length of `self`).
1503 sorter : 1-D array-like, optional
1504 Optional array of integer indices that sort array a into ascending
1505 order. They are typically the result of argsort.
1506
1507 Returns
1508 -------
1509 array of ints or int
1510 If value is array-like, array of insertion points.
1511 If value is scalar, a single integer.
1512
1513 See Also
1514 --------
1515 numpy.searchsorted : Similar method from NumPy.
1516 """
1517 if sorter is not None:
1518 sorter = ensure_platform_int(sorter)
1519
1520 if (
1521 isinstance(arr, np.ndarray)
1522 and arr.dtype.kind in "iu"
1523 and (is_integer(value) or is_integer_dtype(value))
1524 ):
1525 # if `arr` and `value` have different dtypes, `arr` would be
1526 # recast by numpy, causing a slow search.
1527 # Before searching below, we therefore try to give `value` the
1528 # same dtype as `arr`, while guarding against integer overflows.
1529 iinfo = np.iinfo(arr.dtype.type)
1530 value_arr = np.array([value]) if is_integer(value) else np.array(value)
1531 if (value_arr >= iinfo.min).all() and (value_arr <= iinfo.max).all():
1532 # value within bounds, so no overflow, so can convert value dtype
1533 # to dtype of arr
1534 dtype = arr.dtype
1535 else:
1536 dtype = value_arr.dtype
1537
1538 if is_integer(value):
1539 # We know that value is int
1540 value = cast("int", dtype.type(value))
1541 else:
1542 value = pd_array(cast("ArrayLike", value), dtype=dtype)
1543 else:
1544 # E.g. if `arr` is an array with dtype='datetime64[ns]'
1545 # and `value` is a pd.Timestamp, we may need to convert value
1546 arr = ensure_wrapped_if_datetimelike(arr)
1547
1548 # Argument 1 to "searchsorted" of "ndarray" has incompatible type
1549 # "Union[NumpyValueArrayLike, ExtensionArray]"; expected "NumpyValueArrayLike"
1550 return arr.searchsorted(value, side=side, sorter=sorter) # type: ignore[arg-type]
1551
1552
1553# ---- #
1554# diff #
1555# ---- #
1556
1557_diff_special = {"float64", "float32", "int64", "int32", "int16", "int8"}
1558
1559
1560def diff(arr, n: int | float | np.integer | np.floating, axis: AxisInt = 0):
1561 """
1562 difference of n between self,
1563 analogous to s-s.shift(n)
1564
1565 Parameters
1566 ----------
1567 arr : ndarray or ExtensionArray
1568 n : int
1569 number of periods
1570 axis : {0, 1}
1571 axis to shift on
1572
1573 Returns
1574 -------
1575 shifted
1576 """
1577
1578 # added a check on the integer value of period
1579 # see https://github.com/pandas-dev/pandas/issues/56607
1580 if not lib.is_integer(n):
1581 if not (is_float(n) and n.is_integer()):
1582 raise ValueError("periods must be an integer")
1583 n = int(n)
1584 na = np.nan
1585 dtype = arr.dtype
1586
1587 is_bool = is_bool_dtype(dtype)
1588 if is_bool:
1589 op = operator.xor
1590 else:
1591 op = operator.sub
1592
1593 if isinstance(dtype, NumpyEADtype):
1594 # NumpyExtensionArray cannot necessarily hold shifted versions of itself.
1595 arr = arr.to_numpy()
1596 dtype = arr.dtype
1597
1598 if not isinstance(arr, np.ndarray):
1599 # i.e ExtensionArray
1600 if hasattr(arr, f"__{op.__name__}__"):
1601 if axis >= arr.ndim:
1602 raise ValueError(f"cannot diff {type(arr).__name__} on axis={axis}")
1603 return op(arr, arr.shift(n))
1604 else:
1605 raise TypeError(
1606 f"{type(arr).__name__} has no 'diff' method. "
1607 "Convert to a suitable dtype prior to calling 'diff'."
1608 )
1609
1610 is_timedelta = False
1611 if arr.dtype.kind in "mM":
1612 dtype = np.int64
1613 arr = arr.view("i8")
1614 na = iNaT
1615 is_timedelta = True
1616
1617 elif is_bool:
1618 # We have to cast in order to be able to hold np.nan
1619 dtype = np.object_
1620
1621 elif dtype.kind in "iu":
1622 # We have to cast in order to be able to hold np.nan
1623
1624 # int8, int16 are incompatible with float64,
1625 # see https://github.com/cython/cython/issues/2646
1626 if arr.dtype.name in ["int8", "int16"]:
1627 dtype = np.float32
1628 else:
1629 dtype = np.float64
1630
1631 orig_ndim = arr.ndim
1632 if orig_ndim == 1:
1633 # reshape so we can always use algos.diff_2d
1634 arr = arr.reshape(-1, 1)
1635 # TODO: require axis == 0
1636
1637 dtype = np.dtype(dtype)
1638 out_arr = np.empty(arr.shape, dtype=dtype)
1639
1640 na_indexer = [slice(None)] * 2
1641 na_indexer[axis] = slice(None, n) if n >= 0 else slice(n, None)
1642 out_arr[tuple(na_indexer)] = na
1643
1644 if arr.dtype.name in _diff_special:
1645 # TODO: can diff_2d dtype specialization troubles be fixed by defining
1646 # out_arr inside diff_2d?
1647 algos.diff_2d(arr, out_arr, int(n), axis, datetimelike=is_timedelta)
1648 else:
1649 # To keep mypy happy, _res_indexer is a list while res_indexer is
1650 # a tuple, ditto for lag_indexer.
1651 _res_indexer = [slice(None)] * 2
1652 _res_indexer[axis] = slice(n, None) if n >= 0 else slice(None, n)
1653 res_indexer = tuple(_res_indexer)
1654
1655 _lag_indexer = [slice(None)] * 2
1656 _lag_indexer[axis] = slice(None, -n) if n > 0 else slice(-n, None)
1657 lag_indexer = tuple(_lag_indexer)
1658
1659 out_arr[res_indexer] = op(arr[res_indexer], arr[lag_indexer])
1660
1661 if is_timedelta:
1662 out_arr = out_arr.view("timedelta64[ns]")
1663
1664 if orig_ndim == 1:
1665 out_arr = out_arr[:, 0]
1666 return out_arr
1667
1668
1669# --------------------------------------------------------------------
1670# Helper functions
1671
1672
1673# Minimum number of values before safe_sort's counting-sort path is worth it.
1674# Below this, argsort's n*log(n) is cheaper than the linear counting work.
1675_MIN_COUNTING_SORT_LEN = 5_000
1676
1677
1678# Note: safe_sort is in algorithms.py instead of sorting.py because it is
1679# low-dependency, is used in this module, and used private methods from
1680# this module.
1681def safe_sort(
1682 values: Index | ArrayLike,
1683 codes: npt.NDArray[np.intp] | None = None,
1684 use_na_sentinel: bool = True,
1685 assume_unique: bool = False,
1686 verify: bool = True,
1687) -> AnyArrayLike | tuple[AnyArrayLike, np.ndarray]:
1688 """
1689 Sort ``values`` and reorder corresponding ``codes``.
1690
1691 ``values`` should be unique if ``codes`` is not None.
1692 Safe for use with mixed types (int, str), orders ints before strs.
1693
1694 Parameters
1695 ----------
1696 values : list-like
1697 Sequence; must be unique if ``codes`` is not None.
1698 codes : np.ndarray[intp] or None, default None
1699 Indices to ``values``. All out of bound indices are treated as
1700 "not found" and will be masked with ``-1``.
1701 use_na_sentinel : bool, default True
1702 If True, the sentinel -1 will be used for NaN values. If False,
1703 NaN values will be encoded as non-negative integers and will not drop the
1704 NaN from the uniques of the values.
1705 assume_unique : bool, default False
1706 When True, ``values`` are assumed to be unique, which can speed up
1707 the calculation. Ignored when ``codes`` is None.
1708 verify : bool, default True
1709 Check if codes are out of bound for the values and put out of bound
1710 codes equal to ``-1``. If ``verify=False``, it is assumed there
1711 are no out of bound codes. Ignored when ``codes`` is None.
1712
1713 Returns
1714 -------
1715 ordered : AnyArrayLike
1716 Sorted ``values``
1717 new_codes : ndarray
1718 Reordered ``codes``; returned when ``codes`` is not None.
1719
1720 Raises
1721 ------
1722 TypeError
1723 * If ``values`` is not list-like or if ``codes`` is neither None
1724 nor list-like
1725 * If ``values`` cannot be sorted
1726 ValueError
1727 * If ``codes`` is not None and ``values`` contain duplicates.
1728 """
1729 if not isinstance(values, (np.ndarray, ABCExtensionArray, ABCIndex)):
1730 raise TypeError(
1731 "Only np.ndarray, ExtensionArray, and Index objects are allowed to "
1732 "be passed to safe_sort as values"
1733 )
1734
1735 sorter = None
1736 ordered: AnyArrayLike
1737
1738 # For integer values spanning a modest range, the codes remapping below
1739 # can rank each value within that range (a counting sort) instead of
1740 # argsorting.
1741 use_counting = False
1742 vmin = rng_size = 0
1743 if (
1744 codes is not None
1745 and isinstance(values, np.ndarray)
1746 and values.dtype.kind in "iu"
1747 and len(values) >= _MIN_COUNTING_SORT_LEN
1748 ):
1749 vmin = int(values.min())
1750 rng_size = int(values.max()) - vmin + 1
1751 use_counting = rng_size <= 4 * len(values)
1752
1753 if use_counting:
1754 ordered = np.sort(cast("np.ndarray", values))
1755 elif (
1756 not isinstance(values.dtype, ExtensionDtype)
1757 and lib.infer_dtype(values, skipna=False) == "mixed-integer"
1758 ):
1759 ordered = _sort_mixed(values)
1760 else:
1761 try:
1762 if (
1763 codes is None
1764 and isinstance(values, np.ndarray)
1765 and values.dtype.kind in "iufb"
1766 ):
1767 # we only need the sorted values here, and a direct sort
1768 # is cheaper than argsort + take
1769 ordered = np.sort(values)
1770 else:
1771 sorter = values.argsort()
1772 ordered = values.take(sorter)
1773 except (TypeError, decimal.InvalidOperation):
1774 # Previous sorters failed or were not applicable, try `_sort_mixed`
1775 # which would work, but which fails for special case of 1d arrays
1776 # with tuples.
1777 if values.size and isinstance(values[0], tuple):
1778 # error: Argument 1 to "_sort_tuples" has incompatible type
1779 # "Union[Index, ExtensionArray, ndarray[Any, Any]]"; expected
1780 # "ndarray[Any, Any]"
1781 ordered = _sort_tuples(values) # type: ignore[arg-type]
1782 else:
1783 ordered = _sort_mixed(values)
1784
1785 # codes:
1786
1787 if codes is None:
1788 return ordered
1789
1790 if not is_list_like(codes):
1791 raise TypeError(
1792 "Only list-like objects or None are allowed to "
1793 "be passed to safe_sort as codes"
1794 )
1795 codes = ensure_platform_int(np.asarray(codes))
1796
1797 # ranks[i] gives the position of values[i] in `ordered`
1798 if use_counting:
1799 arr = cast("np.ndarray", values)
1800 if arr.dtype.kind == "i":
1801 # go through int64 so differences don't overflow narrower signed
1802 # dtypes; int64 wraparound is exact since the true differences
1803 # are within rng_size
1804 shifted = arr.astype(np.int64, copy=False) - vmin
1805 else:
1806 # unsigned: differences always fit the unsigned dtype
1807 shifted = arr - vmin
1808 present = np.zeros(rng_size, dtype=bool)
1809 present[shifted] = True
1810 counts = present.cumsum(dtype=np.intp)
1811 # the counting pass gives the uniqueness check for free
1812 if not assume_unique and counts[-1] != len(values):
1813 raise ValueError("values should be unique if codes is not None")
1814 ranks = counts[shifted]
1815 ranks -= 1
1816 else:
1817 if not assume_unique and not len(unique(values)) == len(values):
1818 raise ValueError("values should be unique if codes is not None")
1819
1820 if sorter is not None:
1821 # sorter is a permutation, so scatter is a faster equivalent to
1822 # `ranks = sorter.argsort()`
1823 ranks = np.empty(len(sorter), dtype=np.intp)
1824 ranks[sorter] = np.arange(len(sorter), dtype=np.intp)
1825 else:
1826 # mixed types
1827 # error: Argument 1 to "_get_hashtable_algo" has incompatible type
1828 # "Union[Index, ExtensionArray, ndarray[Any, Any]]"; expected
1829 # "ndarray[Any, Any]"
1830 hash_klass, values = _get_hashtable_algo(values) # type: ignore[arg-type]
1831 t = hash_klass(len(values))
1832 # error: Argument 1 to "map_locations" of "HashTable" has incompatible
1833 # type "ExtensionArray | ndarray[Any, Any] | Index | Series";
1834 # expected "ndarray"
1835 t.map_locations(ordered) # type: ignore[arg-type]
1836 ranks = ensure_platform_int(t.lookup(values))
1837
1838 if use_na_sentinel:
1839 # take_nd is faster, but only works for na_sentinels of -1
1840 if verify:
1841 mask = (codes < -len(values)) | (codes >= len(values))
1842 codes[mask] = -1
1843 new_codes = take_nd(ranks, codes, fill_value=-1)
1844 else:
1845 # Out of bound indices will be masked with `-1` next, so we
1846 # may deal with them here without performance loss using `mode='wrap'`
1847 new_codes = ranks.take(codes, mode="wrap")
1848
1849 return ordered, ensure_platform_int(new_codes)
1850
1851
1852def _sort_mixed(values) -> AnyArrayLike:
1853 """order ints before strings before nulls in 1d arrays"""
1854 str_pos = np.array([isinstance(x, str) for x in values], dtype=bool)
1855 null_pos = np.array([isna(x) for x in values], dtype=bool)
1856 num_pos = ~str_pos & ~null_pos
1857 str_argsort = np.argsort(values[str_pos])
1858 num_argsort = np.argsort(values[num_pos])
1859 # convert boolean arrays to positional indices, then order by underlying values
1860 str_locs = str_pos.nonzero()[0].take(str_argsort)
1861 num_locs = num_pos.nonzero()[0].take(num_argsort)
1862 null_locs = null_pos.nonzero()[0]
1863 locs = np.concatenate([num_locs, str_locs, null_locs])
1864 return values.take(locs)
1865
1866
1867def _sort_tuples(values: np.ndarray) -> np.ndarray:
1868 """
1869 Convert array of tuples (1d) to array of arrays (2d).
1870 We need to keep the columns separately as they contain different types and
1871 nans (can't use `np.sort` as it may fail when str and nan are mixed in a
1872 column as types cannot be compared).
1873 """
1874 from pandas.core.internals.construction import to_arrays
1875 from pandas.core.sorting import lexsort_indexer
1876
1877 arrays, _ = to_arrays(values, None)
1878 indexer = lexsort_indexer(arrays, orders=True)
1879 return values[indexer]
1880
1881
1882def union_with_duplicates(
1883 lvals: ArrayLike | Index, rvals: ArrayLike | Index
1884) -> ArrayLike | Index:
1885 """
1886 Extracts the union from lvals and rvals with respect to duplicates and nans in
1887 both arrays.
1888
1889 Parameters
1890 ----------
1891 lvals: np.ndarray or ExtensionArray
1892 left values which is ordered in front.
1893 rvals: np.ndarray or ExtensionArray
1894 right values ordered after lvals.
1895
1896 Returns
1897 -------
1898 np.ndarray or ExtensionArray
1899 Containing the unsorted union of both arrays.
1900
1901 Notes
1902 -----
1903 Caller is responsible for ensuring lvals.dtype == rvals.dtype.
1904 """
1905 from pandas import Series
1906
1907 l_count = value_counts_internal(lvals, dropna=False)
1908 r_count = value_counts_internal(rvals, dropna=False)
1909 l_count, r_count = l_count.align(r_count, fill_value=0)
1910 final_count = np.maximum(l_count.values, r_count.values)
1911 final_count = Series(final_count, index=l_count.index, dtype="int", copy=False)
1912 if isinstance(lvals, ABCMultiIndex) and isinstance(rvals, ABCMultiIndex):
1913 unique_vals = lvals.append(rvals).unique()
1914 else:
1915 if isinstance(lvals, ABCIndex):
1916 lvals = lvals._values
1917 if isinstance(rvals, ABCIndex):
1918 rvals = rvals._values
1919 # error: List item 0 has incompatible type "Union[ExtensionArray,
1920 # ndarray[Any, Any], Index]"; expected "Union[ExtensionArray,
1921 # ndarray[Any, Any]]"
1922 combined = concat_compat([lvals, rvals]) # type: ignore[list-item]
1923 unique_vals = unique(combined)
1924 unique_vals = ensure_wrapped_if_datetimelike(unique_vals)
1925 repeats = final_count.reindex(unique_vals).values
1926 return np.repeat(unique_vals, repeats)
1927
1928
1929def map_array(
1930 arr: ArrayLike,
1931 mapper,
1932 na_action: Literal["ignore"] | None = None,
1933) -> np.ndarray | ExtensionArray | Index:
1934 """
1935 Map values using an input mapping or function.
1936
1937 Parameters
1938 ----------
1939 mapper : function, dict, or Series
1940 Mapping correspondence.
1941 na_action : {None, 'ignore'}, default None
1942 If 'ignore', propagate NA values, without passing them to the
1943 mapping correspondence.
1944
1945 Returns
1946 -------
1947 Union[ndarray, Index, ExtensionArray]
1948 The output of the mapping function applied to the array.
1949 If the function returns a tuple with more than one element
1950 a MultiIndex will be returned.
1951 """
1952 from pandas import Index
1953
1954 if na_action not in (None, "ignore"):
1955 msg = f"na_action must either be 'ignore' or None, {na_action} was passed"
1956 raise ValueError(msg)
1957
1958 # we can fastpath dict/Series to an efficient map
1959 # as we know that we are not going to have to yield
1960 # python types
1961 if is_dict_like(mapper):
1962 if isinstance(mapper, dict) and hasattr(mapper, "__missing__"):
1963 # If a dictionary subclass defines a default value method,
1964 # convert mapper to a lookup function (GH #15999).
1965 dict_with_default = mapper
1966 mapper = lambda x: dict_with_default[
1967 np.nan if isinstance(x, float) and np.isnan(x) else x
1968 ]
1969 else:
1970 # Dictionary does not have a default. Thus it's safe to
1971 # convert to a Series for efficiency.
1972 # we specify the keys here to handle the
1973 # possibility that they are tuples
1974
1975 # The return value of mapping with an empty mapper is
1976 # expected to be pd.Series(np.nan, ...). As np.nan is
1977 # of dtype float64 the return value of this method should
1978 # be float64 as well
1979 from pandas import Series
1980
1981 if len(mapper) == 0:
1982 mapper = Series(mapper, dtype=np.float64)
1983 elif isinstance(mapper, dict):
1984 mapper = Series(
1985 mapper.values(), index=Index(mapper.keys(), tupleize_cols=False)
1986 )
1987 else:
1988 mapper = Series(mapper)
1989
1990 if isinstance(mapper, ABCSeries):
1991 if na_action == "ignore":
1992 mapper = mapper[mapper.index.notna()]
1993
1994 # Since values were input this means we came from either
1995 # a dict or a series and mapper should be an index
1996 indexer = mapper.index.get_indexer(arr)
1997 new_values = take_nd(mapper._values, indexer)
1998
1999 return new_values
2000
2001 if not len(arr):
2002 return arr.copy()
2003
2004 # we must convert to python types
2005 values = arr.astype(object, copy=False)
2006 if na_action is None:
2007 return lib.map_infer(values, mapper)
2008 else:
2009 return lib.map_infer_mask(values, mapper, mask=isna(values).view(np.uint8))